Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ processing \ aurora \ channel_nomenclature.py: 96%
109 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-10 00:11 -0800
1# =====================================================
2# Imports
3# =====================================================
4from typing import Annotated
6from pydantic import computed_field, Field, field_validator, ValidationInfo
8from mt_metadata.base import MetadataBase
9from mt_metadata.common.enumerations import StrEnumerationBase
10from mt_metadata.transfer_functions import CHANNEL_MAPS
13# =====================================================
14class ExEnum(StrEnumerationBase):
15 ex = "ex"
16 e1 = "e1"
17 e2 = "e2"
18 e3 = "e3"
19 e4 = "e4"
22class EyEnum(StrEnumerationBase):
23 ey = "ey"
24 e1 = "e1"
25 e2 = "e2"
26 e3 = "e3"
27 e4 = "e4"
30class HxEnum(StrEnumerationBase):
31 bx = "bx"
32 hx = "hx"
33 h1 = "h1"
34 h2 = "h2"
35 h3 = "h3"
38class HyEnum(StrEnumerationBase):
39 by = "by"
40 hy = "hy"
41 h1 = "h1"
42 h2 = "h2"
43 h3 = "h3"
46class HzEnum(StrEnumerationBase):
47 bz = "bz"
48 hz = "hz"
49 h1 = "h1"
50 h2 = "h2"
51 h3 = "h3"
54class SupportedNomenclatureEnum(StrEnumerationBase):
55 default = "default"
56 lemi12 = "lemi12"
57 lemi34 = "lemi34"
58 musgraves = "musgraves"
59 phoenix123 = "phoenix123"
62class ChannelNomenclature(MetadataBase):
63 ex: Annotated[
64 ExEnum,
65 Field(
66 default="ex",
67 description="label for the X electric field channel, X is assumed to be North",
68 alias=None,
69 json_schema_extra={
70 "units": None,
71 "required": True,
72 "examples": ["ex"],
73 },
74 ),
75 ]
77 ey: Annotated[
78 EyEnum,
79 Field(
80 default="ey",
81 description="label for the Y electric field channel, Y is assumed to be East",
82 alias=None,
83 json_schema_extra={
84 "units": None,
85 "required": True,
86 "examples": ["ey"],
87 },
88 ),
89 ]
91 hx: Annotated[
92 HxEnum,
93 Field(
94 default="hx",
95 description="label for the X magnetic field channel, X is assumed to be North",
96 alias=None,
97 json_schema_extra={
98 "units": None,
99 "required": True,
100 "examples": ["hx"],
101 },
102 ),
103 ]
105 hy: Annotated[
106 HyEnum,
107 Field(
108 default="hy",
109 description="label for the Y magnetic field channel, Y is assumed to be East",
110 alias=None,
111 json_schema_extra={
112 "units": None,
113 "required": True,
114 "examples": ["hy"],
115 },
116 ),
117 ]
119 hz: Annotated[
120 HzEnum,
121 Field(
122 default="hz",
123 description="label for the Z magnetic field channel, Z is assumed to be vertical Down",
124 alias=None,
125 json_schema_extra={
126 "units": None,
127 "required": True,
128 "examples": ["hz"],
129 },
130 ),
131 ]
133 keyword: Annotated[
134 SupportedNomenclatureEnum,
135 Field(
136 default=SupportedNomenclatureEnum.default,
137 description="Keyword for the channel nomenclature system",
138 alias=None,
139 json_schema_extra={
140 "units": None,
141 "required": False,
142 "examples": ["default", "lemi12", "lemi34", "musgraves", "phoenix123"],
143 },
144 ),
145 ]
147 @field_validator("keyword", mode="before")
148 @classmethod
149 def check_keyword(cls, value, info: ValidationInfo):
150 if value is None:
151 value = "default"
152 return value
154 @computed_field
155 @property
156 def ex_ey(self) -> list[str]:
157 return [self.ex.value, self.ey.value]
159 @computed_field
160 @property
161 def hx_hy(self) -> list[str]:
162 return [self.hx.value, self.hy.value]
164 @computed_field
165 @property
166 def hx_hy_hz(self) -> list[str]:
167 return [self.hx.value, self.hy.value, self.hz.value]
169 @computed_field
170 @property
171 def ex_ey_hz(self) -> list[str]:
172 return [self.ex.value, self.ey.value, self.hz.value]
174 @computed_field
175 @property
176 def default_input_channels(self) -> list[str]:
177 return self.hx_hy
179 @computed_field
180 @property
181 def default_output_channels(self) -> list[str]:
182 return self.ex_ey_hz
184 @computed_field
185 @property
186 def default_reference_channels(self) -> list[str]:
187 return self.hx_hy
189 def get_channel_map(self) -> dict[str, str]:
190 """
191 Based on self.keyword return the mapping between conventional channel names and
192 the custom channel names in the particular nomenclature.
194 """
195 try:
196 return CHANNEL_MAPS[self.keyword.lower()]
197 except KeyError:
198 msg = f"channel mt_system {self.keyword} unknown)"
199 raise NotImplementedError(msg)
201 def update(self) -> None:
202 """
203 Assign values to standard channel names "ex", "ey" etc based on channel_map dict
204 """
205 channel_map = self.get_channel_map()
206 self.ex = channel_map["ex"] # type: ignore
207 self.ey = channel_map["ey"] # type: ignore
208 self.hx = channel_map["hx"] # type: ignore
209 self.hy = channel_map["hy"] # type: ignore
210 self.hz = channel_map["hz"] # type: ignore
212 def unpack(self) -> tuple:
213 return self.ex.value, self.ey.value, self.hx.value, self.hy.value, self.hz.value
215 @computed_field
216 @property
217 def channels(self) -> list[str]:
218 channels = list(self.get_channel_map().values())
219 return channels
221 def __setattr__(self, name, value):
222 """Override setattr to automatically update channels when keyword changes."""
223 # Call parent setattr first
224 super().__setattr__(name, value)
226 # If keyword was changed and this is not during initial construction,
227 # update the channel mappings
228 if (
229 name == "keyword"
230 and hasattr(self, "_initialized")
231 and getattr(self, "_initialized", False)
232 ):
233 self.update()
235 def model_post_init(self, __context):
236 """Called after model initialization to set up auto-update and do initial update."""
237 self._initialized = True
238 self.update()