Source code for uwacan.background

 1"""Classes and methods to evaluate, model, and compensate for background noise in measurements.
 2
 3.. autosummary::
 4    :toctree: generated
 5
 6    Background
 7
 8"""
 9
10from . import _core
11import xarray as xr
12import numpy as np
13
14

[docs]
15class Background(_core.FrequencyData):
16    """A class for simple measured background noise.
17
18    Parameters
19    ----------
20    data : `~uwacan.FrequencyData`
21        The measured background noise, as a a power spectral density.
22    snr_requirement : float
23        The required SnR for a measurement to be valid.
24        The compensation will output NaN for invalid
25        data points.
26    """
27
28    def __init__(self, data, snr_requirement=3, **kwargs):
29        super().__init__(data, **kwargs)
30        self.snr_requirement = snr_requirement
31

[docs]
32    def __call__(self, sensor_power):
33        """Compensate a recorded power spectral density.
34
35        Parameters
36        ----------
37        sensor_power : `~uwacan.FrequencyData`
38            The measured power spectral density to compensate.
39            The background measurement will be interpolated
40            to the required frequencies if needed.
41
42        Notes
43        -----
44        We have requirements on the sensor information on the
45        background data and the sensor data.
46
47        1) If the background data has sensor information,
48           the recorded power also needs to have sensor data.
49        2) If the background data has no sensor information, it does
50           not matter if the recorded power has sensor information.
51        3) If both have sensor information, all the sensors in the
52           recorded power has to exist in the background data.
53
54        """
55        background = self.data
56        sensor_power_xr = sensor_power.data if isinstance(sensor_power, _core.xrwrap) else sensor_power
57
58        # if bg has sensors, data needs to have at least the same sensors
59        if "sensor" in background.coords:
60            if "sensor" not in sensor_power_xr.coords:
61                raise ValueError("Cannot apply sensor-wise background compensation to sensor-less recording")
62            if "sensor" not in background.dims:
63                # Single sensor in background, expand it to a dim so we can select from it
64                background = background.expand_dims("sensor")
65            # Pick the correct sensors from the background
66            background = background.sel(sensor=sensor_power_xr.coords["sensor"])
67
68        if not sensor_power_xr.frequency.equals(background.frequency):
69            background_interp = background.interp(
70                frequency=sensor_power_xr.frequency,
71                method="linear",
72            )
73            # Extrapolating using the lowest and highest frequency in the background
74            background_interp = xr.where(
75                background_interp.frequency <= background.frequency[0],
76                background.isel(frequency=0),
77                background_interp,
78            )
79            background_interp = xr.where(
80                background_interp.frequency >= background.frequency[-1],
81                background.isel(frequency=-1),
82                background_interp,
83            )
84            background = background_interp
85
86        snr = _core.dB(sensor_power_xr / background, power=True)
87        compensated = xr.where(
88            snr > self.snr_requirement,
89            sensor_power_xr - background,
90            np.nan,
91        )
92        compensated = sensor_power.from_dataset(compensated)
93        compensated.snr = snr
94        return compensated