modified: checkpoints_final/task1_final/task1_final_model.pkl new file: checkpoints_final/task1_final/task1_tmd_adapter.pkl new file: scripts_4/README.md new file: scripts_4/__init__.py new file: scripts_4/__pycache__/adapter.cpython-314.pyc new file: scripts_4/adapter.py new file: scripts_4/predict_tmd_single.py new file: scripts_4/train_tmd_adapter.py
102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import pickle
|
|
|
|
import numpy as np
|
|
|
|
|
|
def _safe_float(value: float) -> float:
|
|
return float(np.asarray(value, dtype=np.float64).reshape(-1)[0])
|
|
|
|
|
|
@dataclass
|
|
class FrequencyScaleAdapter:
|
|
frequencies_hz: np.ndarray
|
|
scale_values: np.ndarray
|
|
shrinkage: float = 0.2
|
|
extrapolation_decay_hz: float = 0.25
|
|
min_scale: float = 0.4
|
|
max_scale: float = 2.5
|
|
|
|
def __post_init__(self) -> None:
|
|
freqs = np.asarray(self.frequencies_hz, dtype=np.float64).reshape(-1)
|
|
scales = np.asarray(self.scale_values, dtype=np.float64).reshape(-1)
|
|
if freqs.size == 0:
|
|
raise ValueError("frequencies_hz cannot be empty.")
|
|
if freqs.size != scales.size:
|
|
raise ValueError("frequencies_hz and scale_values must have the same length.")
|
|
order = np.argsort(freqs)
|
|
freqs = freqs[order]
|
|
scales = np.clip(scales[order], self.min_scale, self.max_scale)
|
|
# Pull scales toward 1.0 to reduce overfitting with tiny TMD samples.
|
|
scales = 1.0 + (scales - 1.0) * (1.0 - float(np.clip(self.shrinkage, 0.0, 0.95)))
|
|
self.frequencies_hz = freqs
|
|
self.scale_values = scales
|
|
|
|
def _inside_range_weight(self, frequency_hz: float) -> float:
|
|
left = self.frequencies_hz[0]
|
|
right = self.frequencies_hz[-1]
|
|
f = _safe_float(frequency_hz)
|
|
if left <= f <= right:
|
|
return 1.0
|
|
distance = min(abs(f - left), abs(f - right))
|
|
decay = max(_safe_float(self.extrapolation_decay_hz), 1e-6)
|
|
return float(np.exp(-distance / decay))
|
|
|
|
def scale_at(self, frequency_hz: float) -> float:
|
|
f = _safe_float(frequency_hz)
|
|
boundary_interp = float(np.interp(f, self.frequencies_hz, self.scale_values))
|
|
in_range_weight = self._inside_range_weight(f)
|
|
scale = 1.0 + in_range_weight * (boundary_interp - 1.0)
|
|
return float(np.clip(scale, self.min_scale, self.max_scale))
|
|
|
|
def predict(self, base_rms: float, frequency_hz: float) -> float:
|
|
return max(_safe_float(base_rms), 0.0) * self.scale_at(frequency_hz)
|
|
|
|
def to_payload(self) -> dict:
|
|
return {
|
|
"frequencies_hz": self.frequencies_hz.tolist(),
|
|
"scale_values": self.scale_values.tolist(),
|
|
# scale_values are already post-shrinkage values after __post_init__.
|
|
"shrinkage": 0.0,
|
|
"extrapolation_decay_hz": float(self.extrapolation_decay_hz),
|
|
"min_scale": float(self.min_scale),
|
|
"max_scale": float(self.max_scale),
|
|
"already_shrunk": True,
|
|
}
|
|
|
|
@classmethod
|
|
def from_payload(cls, payload: dict) -> "FrequencyScaleAdapter":
|
|
shrinkage = float(payload.get("shrinkage", 0.2))
|
|
if bool(payload.get("already_shrunk", False)):
|
|
shrinkage = 0.0
|
|
return cls(
|
|
frequencies_hz=np.asarray(payload["frequencies_hz"], dtype=np.float64),
|
|
scale_values=np.asarray(payload["scale_values"], dtype=np.float64),
|
|
shrinkage=shrinkage,
|
|
extrapolation_decay_hz=float(payload.get("extrapolation_decay_hz", 0.25)),
|
|
min_scale=float(payload.get("min_scale", 0.4)),
|
|
max_scale=float(payload.get("max_scale", 2.5)),
|
|
)
|
|
|
|
|
|
def save_adapter(adapter: FrequencyScaleAdapter, output_path: Path, extra: dict | None = None) -> None:
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = {
|
|
"adapter_type": "frequency_scale_interp_v1",
|
|
"adapter": adapter.to_payload(),
|
|
"extra": extra or {},
|
|
}
|
|
with output_path.open("wb") as handle:
|
|
pickle.dump(payload, handle)
|
|
|
|
|
|
def load_adapter(adapter_path: Path) -> tuple[FrequencyScaleAdapter, dict]:
|
|
with adapter_path.open("rb") as handle:
|
|
payload = pickle.load(handle)
|
|
adapter = FrequencyScaleAdapter.from_payload(payload["adapter"])
|
|
extra = payload.get("extra", {})
|
|
return adapter, extra
|