Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ common \ band.py: 97%
117 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-10 00:11 -0800
1"""
2Band class for frequency band definitions.
4Development Notes:
5 To add better overlap and intersection checking, consider using piso
6 https://piso.readthedocs.io/en/latest/getting_started/index.html
7"""
9# =====================================================
10# Imports
11# =====================================================
12from typing import Annotated, Optional
14import numpy as np
15import pandas as pd
16from pydantic import (
17 computed_field,
18 Field,
19 field_validator,
20 model_validator,
21 ValidationInfo,
22)
24from mt_metadata.base import MetadataBase
25from mt_metadata.common.enumerations import StrEnumerationBase
28# =====================================================
29class CenterAveragingTypeEnum(StrEnumerationBase):
30 arithmetic = "arithmetic"
31 geometric = "geometric"
34class ClosedEnum(StrEnumerationBase):
35 left = "left"
36 right = "right"
37 both = "both"
40class Band(MetadataBase):
41 decimation_level: Annotated[
42 int,
43 Field(
44 default=None,
45 description="Decimation level for the band",
46 alias=None,
47 json_schema_extra={
48 "units": None,
49 "required": True,
50 "examples": ["0"],
51 },
52 ),
53 ]
55 index_max: Annotated[
56 int,
57 Field(
58 default=None,
59 description="maximum band index",
60 alias=None,
61 json_schema_extra={
62 "units": None,
63 "required": True,
64 "examples": ["10"],
65 },
66 ),
67 ]
69 index_min: Annotated[
70 int,
71 Field(
72 default=None,
73 description="minimum band index",
74 alias=None,
75 json_schema_extra={
76 "units": None,
77 "required": True,
78 "examples": ["10"],
79 },
80 ),
81 ]
83 frequency_max: Annotated[
84 float,
85 Field(
86 default=0.0,
87 description="maximum band frequency",
88 alias=None,
89 json_schema_extra={
90 "units": "Hertz",
91 "required": True,
92 "examples": ["0.04296875"],
93 },
94 ),
95 ]
97 frequency_min: Annotated[
98 float,
99 Field(
100 default=0.0,
101 description="minimum band frequency",
102 alias=None,
103 json_schema_extra={
104 "units": "Hertz",
105 "required": True,
106 "examples": ["0.03515625"],
107 },
108 ),
109 ]
111 center_averaging_type: Annotated[
112 CenterAveragingTypeEnum,
113 Field(
114 default=CenterAveragingTypeEnum.geometric,
115 description="type of average to apply when computing the band center",
116 alias=None,
117 json_schema_extra={
118 "units": None,
119 "required": True,
120 "examples": ["geometric"],
121 },
122 ),
123 ]
125 closed: Annotated[
126 ClosedEnum,
127 Field(
128 default=ClosedEnum.left,
129 description="whether interval is open or closed",
130 alias=None,
131 json_schema_extra={
132 "units": None,
133 "required": True,
134 "examples": ["left"],
135 },
136 ),
137 ]
139 name: Annotated[
140 Optional[str],
141 Field(
142 default="",
143 description="Name of the band",
144 alias=None,
145 json_schema_extra={
146 "units": None,
147 "required": False,
148 "examples": ["0.039062"],
149 },
150 ),
151 ]
153 # we want a name to be generated if not provided, but we also want to allow
154 # the user to set a name explicitly, so we use a field validator that runs after,
155 # but we also need a model validator for after.
156 @field_validator("name", mode="after")
157 @classmethod
158 def validate_name(cls, value: Optional[str], info: ValidationInfo) -> str:
159 if value in ["", None]:
160 # Generate a default name using available data
161 if "frequency_min" in info.data and "frequency_max" in info.data:
162 center_freq = (
163 info.data["frequency_min"] + info.data["frequency_max"]
164 ) / 2
165 return f"{center_freq:.6f}"
166 else:
167 return "unnamed_band"
168 elif not isinstance(value, str):
169 raise TypeError(f"Expected string, got {type(value)}")
170 else:
171 return value
173 @field_validator("frequency_min", "frequency_max", mode="after")
174 @classmethod
175 def update_name_on_frequency_change(
176 cls, value: float, info: ValidationInfo
177 ) -> float:
178 # This will trigger a model validation after the field is set
179 return value
181 @model_validator(mode="after")
182 def check_name(self) -> "Band":
183 # Update name if it's the default "empty_band" and we now have frequencies
184 if self.name in ["", None] or (
185 self.name == "empty_band"
186 and self.frequency_min != 0.0
187 and self.frequency_max != 0.0
188 ):
189 if self.frequency_min != 0.0 and self.frequency_max != 0.0:
190 center_freq = (self.frequency_min + self.frequency_max) / 2
191 self.name = f"{center_freq:.6f}"
192 else:
193 self.name = "empty_band"
194 return self
196 @computed_field
197 @property
198 def lower_bound(self) -> float:
199 return self.frequency_min
201 @computed_field
202 @property
203 def upper_bound(self) -> float:
204 return self.frequency_max
206 @computed_field
207 @property
208 def width(self) -> float:
209 """returns the width of the band (the bandwidth)."""
210 return self.upper_bound - self.lower_bound
212 @computed_field
213 @property
214 def lower_closed(self) -> bool:
215 return self.to_interval().closed_left
217 @computed_field
218 @property
219 def upper_closed(self) -> bool:
220 return self.to_interval().closed_right
222 def _indices_from_frequencies(self, frequencies: np.ndarray) -> np.ndarray:
223 """
225 Parameters
226 ----------
227 frequencies: numpy array
228 Intended to represent the one-sided (positive) frequency axis of
229 the data that has been FFT-ed
231 Returns
232 -------
233 indices: numpy array of integers
234 Integer indices of the fourier coefficients associated with the
235 frequecies passed as input argument
236 """
237 if self.lower_closed:
238 cond1 = frequencies >= self.lower_bound
239 else:
240 cond1 = frequencies > self.lower_bound
241 if self.upper_closed:
242 cond2 = frequencies <= self.upper_bound
243 else:
244 cond2 = frequencies < self.upper_bound
246 indices = np.where(cond1 & cond2)[0]
247 return indices
249 def set_indices_from_frequencies(self, frequencies: np.ndarray) -> None:
250 """assumes min/max freqs are defined"""
251 indices = self._indices_from_frequencies(frequencies)
252 self.index_min = indices[0]
253 self.index_max = indices[-1]
255 def to_interval(self):
256 # Handle both string and enum values for closed
257 closed_value = (
258 self.closed.value if hasattr(self.closed, "value") else self.closed
259 )
260 return pd.Interval(self.frequency_min, self.frequency_max, closed=closed_value)
262 @property
263 def harmonic_indices(self):
264 """
265 Assumes all harmoincs between min and max are present in the band
267 Returns
268 -------
269 numpy array of integers corresponding to harminic indices
270 """
271 return np.arange(self.index_min, self.index_max + 1)
273 def in_band_harmonics(self, frequencies: np.ndarray):
274 """
275 Parameters
276 ----------
277 frequencies: array-like, floating poirt
279 Returns: numpy array
280 the actual harmonics or frequencies in band, rather than the indices.
281 -------
283 """
284 indices = self._indices_from_frequencies(frequencies)
285 harmonics = frequencies[indices]
286 return harmonics
288 @property
289 def center_frequency(self) -> float:
290 """
291 Returns
292 -------
293 center_frequency: float
294 The frequency associated with the band center.
295 """
296 if self.center_averaging_type == "geometric":
297 return np.sqrt(self.lower_bound * self.upper_bound)
298 elif self.center_averaging_type == "arithmetic":
299 return (self.lower_bound + self.upper_bound) / 2
300 else:
301 # Default fallback, could raise an error or return a default value
302 return float("nan")
304 @property
305 def center_period(self) -> float:
306 """Returns the inverse of center frequency."""
307 return 1.0 / self.center_frequency
309 def overlaps(self, other) -> bool:
310 """Check if this band overlaps with another"""
311 ivl = self.to_interval()
312 other_ivl = other.to_interval()
313 return ivl.overlaps(other_ivl)
315 def contains(self, other) -> bool:
316 """Check if this band contains nother"""
317 ivl = self.to_interval()
318 cond1 = ivl.__contains__(other.lower_bound)
319 cond2 = ivl.__contains__(other.upper_bound)
320 return cond1 & cond2
322 @computed_field
323 @property
324 def fractional_bandwidth(self) -> float:
325 """
326 See
327 - https://en.wikipedia.org/wiki/Bandwidth_(signal_processing)#Fractional_bandwidth
328 - https://en.wikipedia.org/wiki/Q_factor
330 Returns
331 -------
333 """
334 return self.width / self.center_frequency
336 @computed_field
337 @property
338 def Q(self) -> float:
339 """
340 Quality factor (Q) of the band.
342 Returns
343 -------
344 float
345 Q factor. Returns infinity for zero-width bands.
346 """
347 if self.fractional_bandwidth == 0.0:
348 return float("inf")
349 return 1.0 / self.fractional_bandwidth