13365054372

Committed 17 Feb 2025 07:28AM UTC coverage: 73.985% (+1.7%) from 72.288%

Build # 13365054372

Build Type

push

github

Committed by

avivajpeyi

Commit Message

add covar study

Coverage Stats

249 of 442 branches covered (56.33%)

Branch coverage included in aggregate %.

1 of 3 new or added lines in 1 file covered. (33.33%)

917 of 1134 relevant lines covered (80.86%)

0.81 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

83.67

/src/gap_study_utils/analysis_data.py

from typing import Callable, Dict, List, Optional

import numpy as np

from pywavelet.types.wavelet_bins import compute_bins
from pywavelet.types import FrequencySeries, TimeSeries, Wavelet, WaveletMask
from pywavelet.utils import (
    compute_likelihood,
    evolutionary_psd_from_stationary_psd,
)


from eryn.prior import ProbDistContainer, uniform_dist


from .constants import DT, GAP_RANGES, NF, TMAX, TRUES
from .gaps import GapType, GapWindow
from .utils.noise_curves import (
    noise_curve,
    generate_stationary_noise,
)
from .utils.signal_utils import waveform, compute_snr_dict

from .plotting import plot_analysis_data
from . import logger


class AnalysisData:
    """
    Encapsulates data and methods required for a gap study on time series and
    frequency analysis with wavelets, handling configurations for data sampling,
    PSDs, and gap-windowing, with optional plotting.
    """

    @classmethod
    def DEFAULT(cls):
        data_kwargs = {
            "dt": DT,
            "tmax": TMAX,
            "alpha": 0.0,
            "highpass_fmin": None,
            "frange": None,
            "tgaps": None,
            "noise": False,
            "seed": None,
            "noise_curve": "TDI1",
        }
        gap_kwargs = {
            "type": GapType.STITCH,
            "gap_ranges": GAP_RANGES,
        }
        waveform_generator = waveform
        waveform_parameters = TRUES
        return cls(
            data_kwargs, gap_kwargs, waveform_generator, waveform_parameters
        )

    def __init__(
            self,
            data_kwargs: Optional[Dict] = None,
            gap_kwargs: Optional[Dict] = None,
            waveform_generator: Optional[Callable[..., TimeSeries]] = None,
            waveform_parameters: Optional[List[float]] = None,
            parameter_ranges: Optional[List[List[float]]] = None,
            plotfn: Optional[str] = None,
    ):
        self.data_kwargs = data_kwargs or {}
        self.gap_kwargs = gap_kwargs or {}

        self.waveform_generator = waveform_generator
        self.waveform_parameters = waveform_parameters
        self._initialize_data_params()
        self._initialize_grids()
        self._initialize_gap_window()

        if plotfn:
            _ = self.plot_data(plotfn)
        self.priors:ProbDistContainer = construct_prior(self.waveform_parameters, parameter_ranges)

        logger.info("AnalysisData initialized.")
        logger.info(self.summary)

    def log_prior(self, theta: List[float]) -> float:
        return self.priors.logpdf(theta)

    def _initialize_data_params(self):
        """Initialize core data parameters from `data_kwargs` with default values."""
        self.dt = self.data_kwargs.get("dt", DT)
        self.tmax = self.data_kwargs.get("tmax", TMAX)
        self.alpha = self.data_kwargs.get("alpha", 0.0)
        self.highpass_fmin = self.data_kwargs.get("highpass_fmin", None)  # HARD CODED
        self.frange = self.data_kwargs.get("frange", None)
        self.tgaps = self.data_kwargs.get("tgaps", [])
        if self.tgaps is None:
            self.tgaps = []
        self.noise = self.data_kwargs.get("noise", False)
        self.noise_curve = self.data_kwargs.get("noise_curve", "TDI1")
        self.seed = self.data_kwargs.get("seed", None)
        self.ND = int(self.tmax / self.dt)
        self.time = np.arange(0, self.tmax, self.dt)
        self.freq = np.fft.rfftfreq(self.ND, d=self.dt)
        if self.seed:
            np.random.seed(self.seed)

    def _initialize_grids(self):
        """Compute time and frequency grids for wavelet analysis."""
        self.Nf = self.data_kwargs.get("Nf", NF)
        self.Nt = self.ND // self.Nf
        self.t_grid, self.f_grid = compute_bins(self.Nf, self.Nt, self.tmax)
        # make a mask -- only use f_grid within the frange
        if self.frange is None:
            self.frange = [0, self.freq[-1]]


        # make a mask -- only use f_grid within the frange
        self.mask = WaveletMask.from_restrictions(time_grid=self.t_grid, freq_grid=self.f_grid, frange=self.frange, tgaps=self.tgaps)
        self.zero_wavelet = Wavelet.zeros_from_grid(self.t_grid, self.f_grid)

    def _initialize_gap_window(self):
        """Set up the gap window if `gap_kwargs` are provided."""
        gap_ranges = self.gap_kwargs.get("gap_ranges", [])
        gap_type = self.gap_kwargs.get("type", GapType.STITCH)

        if isinstance(gap_type, str):
            gap_type = GapType[gap_type.upper()]

        if gap_ranges:
            logger.info(f"Initalizing GapWindow with {gap_type} gaps ({len(gap_ranges)} gaps).")
            self.gaps = GapWindow(
                self.time, tmax=self.tmax, gap_ranges=gap_ranges, type=gap_type
            )
        else:
            self.gaps = None

    @property
    def ND(self) -> int:
        return self._ND

    @ND.setter
    def ND(self, n: int):
        """Ensure ND is a power of 2, raising an error with a suggestion if it is not."""
        n_pwr_2 = round_to_nearest_pow_of_2(n)
        if n != n_pwr_2:
            suggested_tmax = n_pwr_2 * self.dt
            raise ValueError(
                f"ND must be a power of 2. "
                f"Given dt={self.dt}, \n"
                f"suggested tmax={suggested_tmax} --> [n=2^{np.log2(n_pwr_2)}],  \n"
                f"current tmax={self.tmax} --> [n=2^{np.log2(n):.2f}]"
            )
        self._ND = n

    @property
    def psd_freqseries(self) -> FrequencySeries:
        """Generate the frequency series PSD if it hasn't been computed."""
        if not hasattr(self, "_psd_freqseries"):
            self._psd_freqseries = noise_curve(self.freq, self.noise_curve)
        return self._psd_freqseries

    @property
    def psd_wavelet(self) -> Wavelet:
        """Compute the wavelet form of the evolutionary PSD."""
        if not hasattr(self, "_psd_wavelet"):
            psd = self.psd_freqseries
            self._psd_wavelet = evolutionary_psd_from_stationary_psd(
                psd.data, psd.freq, self.f_grid, self.t_grid, self.dt
            )
        return self._psd_wavelet

    @property
    def psd_analysis(self) -> Wavelet:
        """Return the PSD for the analysis."""
        if not hasattr(self, "_psd"):
            p = self.psd_wavelet.copy()
            if self.gaps:
                p = self.gaps.apply_nan_gap_to_wavelet(p)
            self._psd = p
        return self._psd

    @property
    def ht(self) -> TimeSeries:
        """Generate the time series from the waveform generator if provided."""
        if not hasattr(self, "_ht"):
            self._ht = (
                self.waveform_generator(*self.waveform_parameters, self.time)
                if self.waveform_generator
                else TimeSeries._EMPTY(self.ND, self.dt)
            )

        return self._ht


    @property
    def chunked_ht(self) -> List[TimeSeries]:
        """Chunk the time series if gaps are present."""
        if not hasattr(self, "_chunked_ht"):
            if self.gaps and self.gaps.type == GapType.STITCH:
                self._chunked_ht = self.gaps._chunk_timeseries(
                    self.ht,alpha=self.alpha, fmin=self.highpass_fmin,
            )
            else:
                self._chunked_ht = [self.ht]
        return self._chunked_ht

    @property
    def chunked_hf(self):
        """Chunk the frequency series if gaps are present."""
        if not hasattr(self, "_chunked_hf"):
            self._chunked_hf = [
                ts.to_frequencyseries() for ts in self.chunked_ht
            ]
        return self._chunked_hf

    @property
    def noise_frequencyseries(self) -> FrequencySeries:
        """Generate stationary noise frequency series if noise is enabled."""
        if not hasattr(self, "_noise_frequencyseries"):
            self._noise_frequencyseries = (
                generate_stationary_noise(
                    ND=self.ND, dt=self.dt, psd=self.psd_freqseries,
                )
                if self.noise
                else FrequencySeries._EMPTY(self.Nf, self.Nt)
            )
        return self._noise_frequencyseries

    @property
    def noise_timeseries(self) -> TimeSeries:
        """Generate stationary noise time series if noise is enabled."""
        if not hasattr(self, "_noise_timeseries"):
            self._noise_timeseries = (
                self.noise_frequencyseries.to_timeseries()
                if self.noise
                else TimeSeries._EMPTY(self.ND, self.dt)
            )
        return self._noise_timeseries

    @property
    def noise_wavelet(self) -> Wavelet:
        """Generate wavelet-transformed noise time series."""
        if not hasattr(self, "_noise_wavelet"):
            if self.noise:
                self._noise_wavelet = self.noise_timeseries.to_wavelet(Nf=self.Nf)
            else:
                self._noise_wavelet = Wavelet.zeros_from_grid(self.t_grid, self.f_grid)
        return self._noise_wavelet

    @property
    def data_timeseries(self) -> TimeSeries:
        """Combine the signal and noise time series."""
        if not hasattr(self, "_data_timeseries"):
            self._data_timeseries = self.ht + self.noise_timeseries
        return self._data_timeseries

    @property
    def hf(self) -> FrequencySeries:
        """Convert time series to frequency series."""
        if not hasattr(self, "_hf"):
            self._hf = self.ht.to_frequencyseries()
        return self._hf

    @property
    def data_frequencyseries(self) -> FrequencySeries:
        """Convert data time series to frequency series."""
        if not hasattr(self, "_data_frequencyseries"):
            self._data_frequencyseries = (
                self.data_timeseries.to_frequencyseries()
            )
        return self._data_frequencyseries

    @property
    def hwavelet(self) -> Wavelet:
        """Compute wavelet transform of the time series."""
        if not hasattr(self, "_hwavelet"):
            self._hwavelet = self.ht.to_wavelet(Nf=self.Nf)
        return self._hwavelet

    @property
    def hwavelet_gapped(self) -> Wavelet:
        """Apply gap windowing to the wavelet-transformed time series."""
        if not hasattr(self, "_hwavelet_gapped"):
            if self.gaps:
                self._hwavelet_gapped = self.gaps.gap_n_transform_timeseries(
                    self.ht, self.Nf, self.alpha, self.highpass_fmin
                )
            else:
                self._hwavelet_gapped = None
        return self._hwavelet_gapped

    @property
    def data_wavelet(self) -> Wavelet:
        """Apply gap windowing and high-pass filtering to data time series and compute wavelet."""
        if not hasattr(self, "_data_wavelet"):
            data_timeseries = self.data_timeseries
            self._data_wavelet = (
                data_timeseries.to_wavelet(Nf=self.Nf)
                if not self.gaps
                else self.gaps.gap_n_transform_timeseries(
                    data_timeseries, self.Nf, self.alpha, self.highpass_fmin
                )
            )
        return self._data_wavelet

    @property
    def summary_dict(self) -> Dict[str, float]:
        """Summary dictionary of analysis metrics, including signal-to-noise ratios (SNR)."""
        if not hasattr(self, "_summary_dict"):
            windowed = self.highpass_fmin is not None and self.highpass_fmin > 0

            self._summary_dict = dict(
                ht=self.ht,
                gaps=self.gaps,
                windowed=windowed,
                noise=self.noise,
                **self.snr_dict,
            )
            self._summary_dict["lnL@true"] = f"{self.lnl(*self.waveform_parameters)}:.2e"

        return self._summary_dict

    @property
    def summary(self) -> str:
        """Formatted summary string of analysis metrics."""
        return "\n".join([f"{k}: {v}" for k, v in self.summary_dict.items()])

    @property
    def snr_dict(self) -> Dict[str, float]:
        """Calculate various SNR values based on the analysis data."""
        if not hasattr(self, "_snr_dict"):
            self._snr_dict = compute_snr_dict(
                self.hf, self.psd_freqseries, self.data_frequencyseries,
                self.hwavelet, self.psd_wavelet, self.data_wavelet,
                self.psd_analysis, self.gaps, self.hwavelet_gapped
            )
        return self._snr_dict



    def plot_data(self, plotfn: str=None, **kwargs):
        """Plot data visualizations including SNR information, time series, wavelet, and PSDs."""
        return plot_analysis_data(self, plotfn, **kwargs)

    def htemplate(self, *args, **kwargs) -> Wavelet:
        ht = self.waveform_generator(*args, **kwargs, t=self.time)
        if self.gaps is not None:
            hwavelet = self.gaps.gap_n_transform_timeseries(
                ht, self.Nf, self.alpha, self.highpass_fmin
            )
        else:
            if self.highpass_fmin:
                ht = ht.highpass_filter(self.highpass_fmin, self.alpha)
            hwavelet = ht.to_wavelet(Nf=self.Nf)
        return hwavelet


    def lnl(self, *args) -> float:
        return compute_likelihood(
            self.data_wavelet, self.htemplate(*args), self.psd_analysis, self.mask
        )

    def noise_lnl(self, *args) -> float:
        return compute_likelihood(
            self.data_wavelet, self.zero_wavelet, self.psd_analysis, self.mask
        )

    def ln_posterior(self, theta:List[float]) -> float:
        if self.log_prior(np.array(theta)) == -np.inf:
            return -np.inf
        return self.lnl(*theta)

    def freqdomain_lnp(self, theta:List[float]) -> float:
        if self.log_prior(np.array(theta)) == -np.inf:
            return -np.inf
        return self.freqdomain_lnl(*theta)


    def freqdomain_lnl(self, *args) -> float:
        ht = self.waveform_generator(*args, t=self.time)
        d = self.data_timeseries
        # if self.highpass_fmin:
        #     ht = ht.highpass_filter(self.highpass_fmin, self.alpha)
        if self.gaps is not None:
            ht.data[self.gaps.gap_bools] = 0
            d.data[self.gaps.gap_bools] = 0
        signal_f = ht.to_frequencyseries().data
        data_f = d.to_frequencyseries().data
        variance_noise_f = (
                self.ND * self.psd_freqseries.data / (4 * self.dt)
        )  # Calculate variance of noise, real and imaginary.
        inn_prod = sum((abs(data_f - signal_f) ** 2) / variance_noise_f)
        return -0.5 * inn_prod



def round_to_nearest_pow_of_2(n: int) -> int:
    return 2 ** int(np.log2(n))


def get_suggested_tmax(tmax:float, dt:float=DT) -> float:
    n = round_to_nearest_pow_of_2(tmax / dt)
    return n * dt


def construct_prior(trues, ranges=None) -> ProbDistContainer:
    ln_a, ln_f, ln_fdot = trues
    if ranges is None:
        lna_range = [ln_a - 0.1, ln_a + 0.1]
        lnf_range = [ln_f - 0.0001, ln_f + 0.0001]
        lnfdot_range = [ln_fdot - 0.0001, ln_fdot + 0.0001]
        ranges = [lna_range, lnf_range, lnfdot_range]
    return ProbDistContainer({i: uniform_dist(*r) for i, r in enumerate(ranges)})



1	from typing import Callable, Dict, List, Optional	1✔
2
3	import numpy as np	1✔
4
5	from pywavelet.types.wavelet_bins import compute_bins	1✔
6	from pywavelet.types import FrequencySeries, TimeSeries, Wavelet, WaveletMask	1✔
7	from pywavelet.utils import (	1✔
8	compute_likelihood,
9	evolutionary_psd_from_stationary_psd,
10	)
11
12
13	from eryn.prior import ProbDistContainer, uniform_dist	1✔
14
15
16	from .constants import DT, GAP_RANGES, NF, TMAX, TRUES	1✔
17	from .gaps import GapType, GapWindow	1✔
18	from .utils.noise_curves import (	1✔
19	noise_curve,
20	generate_stationary_noise,
21	)
22	from .utils.signal_utils import waveform, compute_snr_dict	1✔
23
24	from .plotting import plot_analysis_data	1✔
25	from . import logger	1✔
26
27
28	class AnalysisData:	1✔
29	"""
30	Encapsulates data and methods required for a gap study on time series and
31	frequency analysis with wavelets, handling configurations for data sampling,
32	PSDs, and gap-windowing, with optional plotting.
33	"""
34
35	@classmethod	1✔
36	def DEFAULT(cls):	1✔
37	data_kwargs = {	1✔
38	"dt": DT,
39	"tmax": TMAX,
40	"alpha": 0.0,
41	"highpass_fmin": None,
42	"frange": None,
43	"tgaps": None,
44	"noise": False,
45	"seed": None,
46	"noise_curve": "TDI1",
47	}
48	gap_kwargs = {	1✔
49	"type": GapType.STITCH,
50	"gap_ranges": GAP_RANGES,
51	}
52	waveform_generator = waveform	1✔
53	waveform_parameters = TRUES	1✔
54	return cls(	1✔
55	data_kwargs, gap_kwargs, waveform_generator, waveform_parameters
56	)
57
58	def __init__(	1✔
59	self,
60	data_kwargs: Optional[Dict] = None,
61	gap_kwargs: Optional[Dict] = None,
62	waveform_generator: Optional[Callable[..., TimeSeries]] = None,
63	waveform_parameters: Optional[List[float]] = None,
64	parameter_ranges: Optional[List[List[float]]] = None,
65	plotfn: Optional[str] = None,
66	):
67	self.data_kwargs = data_kwargs or {}	1✔
68	self.gap_kwargs = gap_kwargs or {}	1✔
69
70	self.waveform_generator = waveform_generator	1✔
71	self.waveform_parameters = waveform_parameters	1✔
72	self._initialize_data_params()	1✔
73	self._initialize_grids()	1✔
74	self._initialize_gap_window()	1✔
75
76	if plotfn:	1✔
77	_ = self.plot_data(plotfn)	1✔
78	self.priors:ProbDistContainer = construct_prior(self.waveform_parameters, parameter_ranges)	1✔
79
80	logger.info("AnalysisData initialized.")	1✔
81	logger.info(self.summary)	1✔
82
83	def log_prior(self, theta: List[float]) -> float:	1✔
84	return self.priors.logpdf(theta)	1✔
85
86	def _initialize_data_params(self):	1✔
87	"""Initialize core data parameters from `data_kwargs` with default values."""
88	self.dt = self.data_kwargs.get("dt", DT)	1✔
89	self.tmax = self.data_kwargs.get("tmax", TMAX)	1✔
90	self.alpha = self.data_kwargs.get("alpha", 0.0)	1✔
91	self.highpass_fmin = self.data_kwargs.get("highpass_fmin", None) # HARD CODED	1✔
92	self.frange = self.data_kwargs.get("frange", None)	1✔
93	self.tgaps = self.data_kwargs.get("tgaps", [])	1✔
94	if self.tgaps is None:	1✔
95	self.tgaps = []	1✔
96	self.noise = self.data_kwargs.get("noise", False)	1✔
97	self.noise_curve = self.data_kwargs.get("noise_curve", "TDI1")	1✔
98	self.seed = self.data_kwargs.get("seed", None)	1✔
99	self.ND = int(self.tmax / self.dt)	1✔
100	self.time = np.arange(0, self.tmax, self.dt)	1✔
101	self.freq = np.fft.rfftfreq(self.ND, d=self.dt)	1✔
102	if self.seed:	1!
103	np.random.seed(self.seed)	×
104
105	def _initialize_grids(self):	1✔
106	"""Compute time and frequency grids for wavelet analysis."""
107	self.Nf = self.data_kwargs.get("Nf", NF)	1✔
108	self.Nt = self.ND // self.Nf	1✔
109	self.t_grid, self.f_grid = compute_bins(self.Nf, self.Nt, self.tmax)	1✔
110	# make a mask -- only use f_grid within the frange
111	if self.frange is None:	1!
112	self.frange = [0, self.freq[-1]]	1✔
113
114
115	# make a mask -- only use f_grid within the frange
116	self.mask = WaveletMask.from_restrictions(time_grid=self.t_grid, freq_grid=self.f_grid, frange=self.frange, tgaps=self.tgaps)	1✔
117	self.zero_wavelet = Wavelet.zeros_from_grid(self.t_grid, self.f_grid)	1✔
118
119	def _initialize_gap_window(self):	1✔
120	"""Set up the gap window if `gap_kwargs` are provided."""
121	gap_ranges = self.gap_kwargs.get("gap_ranges", [])	1✔
122	gap_type = self.gap_kwargs.get("type", GapType.STITCH)	1✔
123
124	if isinstance(gap_type, str):	1!
125	gap_type = GapType[gap_type.upper()]	×
126
127	if gap_ranges:	1✔
128	logger.info(f"Initalizing GapWindow with {gap_type} gaps ({len(gap_ranges)} gaps).")	1✔
129	self.gaps = GapWindow(	1✔
130	self.time, tmax=self.tmax, gap_ranges=gap_ranges, type=gap_type
131	)
132	else:
133	self.gaps = None	1✔
134
135	@property	1✔
136	def ND(self) -> int:	1✔
137	return self._ND	1✔
138
139	@ND.setter	1✔
140	def ND(self, n: int):	1✔
141	"""Ensure ND is a power of 2, raising an error with a suggestion if it is not."""
142	n_pwr_2 = round_to_nearest_pow_of_2(n)	1✔
143	if n != n_pwr_2:	1!
144	suggested_tmax = n_pwr_2 * self.dt	×
145	raise ValueError(	×
146	f"ND must be a power of 2. "
147	f"Given dt={self.dt}, \n"
148	f"suggested tmax={suggested_tmax} --> [n=2^{np.log2(n_pwr_2)}], \n"
149	f"current tmax={self.tmax} --> [n=2^{np.log2(n):.2f}]"
150	)
151	self._ND = n	1✔
152
153	@property	1✔
154	def psd_freqseries(self) -> FrequencySeries:	1✔
155	"""Generate the frequency series PSD if it hasn't been computed."""
156	if not hasattr(self, "_psd_freqseries"):	1✔
157	self._psd_freqseries = noise_curve(self.freq, self.noise_curve)	1✔
158	return self._psd_freqseries	1✔
159
160	@property	1✔
161	def psd_wavelet(self) -> Wavelet:	1✔
162	"""Compute the wavelet form of the evolutionary PSD."""
163	if not hasattr(self, "_psd_wavelet"):	1✔
164	psd = self.psd_freqseries	1✔
165	self._psd_wavelet = evolutionary_psd_from_stationary_psd(	1✔
166	psd.data, psd.freq, self.f_grid, self.t_grid, self.dt
167	)
168	return self._psd_wavelet	1✔
169
170	@property	1✔
171	def psd_analysis(self) -> Wavelet:	1✔
172	"""Return the PSD for the analysis."""
173	if not hasattr(self, "_psd"):	1✔
174	p = self.psd_wavelet.copy()	1✔
175	if self.gaps:	1✔
176	p = self.gaps.apply_nan_gap_to_wavelet(p)	1✔
177	self._psd = p	1✔
178	return self._psd	1✔
179
180	@property	1✔
181	def ht(self) -> TimeSeries:	1✔
182	"""Generate the time series from the waveform generator if provided."""
183	if not hasattr(self, "_ht"):	1✔
184	self._ht = (	1✔
185	self.waveform_generator(*self.waveform_parameters, self.time)
186	if self.waveform_generator
187	else TimeSeries._EMPTY(self.ND, self.dt)
188	)
189
190	return self._ht	1✔
191
192
193	@property	1✔
194	def chunked_ht(self) -> List[TimeSeries]:	1✔
195	"""Chunk the time series if gaps are present."""
196	if not hasattr(self, "_chunked_ht"):	1✔
197	if self.gaps and self.gaps.type == GapType.STITCH:	1✔
198	self._chunked_ht = self.gaps._chunk_timeseries(	1✔
199	self.ht,alpha=self.alpha, fmin=self.highpass_fmin,
200	)
201	else:
202	self._chunked_ht = [self.ht]	1✔
203	return self._chunked_ht	1✔
204
205	@property	1✔
206	def chunked_hf(self):	1✔
207	"""Chunk the frequency series if gaps are present."""
208	if not hasattr(self, "_chunked_hf"):	1!
209	self._chunked_hf = [	1✔
210	ts.to_frequencyseries() for ts in self.chunked_ht
211	]
212	return self._chunked_hf	1✔
213
214	@property	1✔
215	def noise_frequencyseries(self) -> FrequencySeries:	1✔
216	"""Generate stationary noise frequency series if noise is enabled."""
217	if not hasattr(self, "_noise_frequencyseries"):	1!
218	self._noise_frequencyseries = (	1✔
219	generate_stationary_noise(
220	ND=self.ND, dt=self.dt, psd=self.psd_freqseries,
221	)
222	if self.noise
223	else FrequencySeries._EMPTY(self.Nf, self.Nt)
224	)
225	return self._noise_frequencyseries	1✔
226
227	@property	1✔
228	def noise_timeseries(self) -> TimeSeries:	1✔
229	"""Generate stationary noise time series if noise is enabled."""
230	if not hasattr(self, "_noise_timeseries"):	1!
231	self._noise_timeseries = (	1✔
232	self.noise_frequencyseries.to_timeseries()
233	if self.noise
234	else TimeSeries._EMPTY(self.ND, self.dt)
235	)
236	return self._noise_timeseries	1✔
237
238	@property	1✔
239	def noise_wavelet(self) -> Wavelet:	1✔
240	"""Generate wavelet-transformed noise time series."""
241	if not hasattr(self, "_noise_wavelet"):	×
242	if self.noise:	×
243	self._noise_wavelet = self.noise_timeseries.to_wavelet(Nf=self.Nf)	×
244	else:
245	self._noise_wavelet = Wavelet.zeros_from_grid(self.t_grid, self.f_grid)	×
246	return self._noise_wavelet	×
247
248	@property	1✔
249	def data_timeseries(self) -> TimeSeries:	1✔
250	"""Combine the signal and noise time series."""
251	if not hasattr(self, "_data_timeseries"):	1✔
252	self._data_timeseries = self.ht + self.noise_timeseries	1✔
253	return self._data_timeseries	1✔
254
255	@property	1✔
256	def hf(self) -> FrequencySeries:	1✔
257	"""Convert time series to frequency series."""
258	if not hasattr(self, "_hf"):	1✔
259	self._hf = self.ht.to_frequencyseries()	1✔
260	return self._hf	1✔
261
262	@property	1✔
263	def data_frequencyseries(self) -> FrequencySeries:	1✔
264	"""Convert data time series to frequency series."""
265	if not hasattr(self, "_data_frequencyseries"):	1!
266	self._data_frequencyseries = (	1✔
267	self.data_timeseries.to_frequencyseries()
268	)
269	return self._data_frequencyseries	1✔
270
271	@property	1✔
272	def hwavelet(self) -> Wavelet:	1✔
273	"""Compute wavelet transform of the time series."""
274	if not hasattr(self, "_hwavelet"):	1✔
275	self._hwavelet = self.ht.to_wavelet(Nf=self.Nf)	1✔
276	return self._hwavelet	1✔
277
278	@property	1✔
279	def hwavelet_gapped(self) -> Wavelet:	1✔
280	"""Apply gap windowing to the wavelet-transformed time series."""
281	if not hasattr(self, "_hwavelet_gapped"):	1✔
282	if self.gaps:	1✔
283	self._hwavelet_gapped = self.gaps.gap_n_transform_timeseries(	1✔
284	self.ht, self.Nf, self.alpha, self.highpass_fmin
285	)
286	else:
287	self._hwavelet_gapped = None	1✔
288	return self._hwavelet_gapped	1✔
289
290	@property	1✔
291	def data_wavelet(self) -> Wavelet:	1✔
292	"""Apply gap windowing and high-pass filtering to data time series and compute wavelet."""
293	if not hasattr(self, "_data_wavelet"):	1✔
294	data_timeseries = self.data_timeseries	1✔
295	self._data_wavelet = (	1✔
296	data_timeseries.to_wavelet(Nf=self.Nf)
297	if not self.gaps
298	else self.gaps.gap_n_transform_timeseries(
299	data_timeseries, self.Nf, self.alpha, self.highpass_fmin
300	)
301	)
302	return self._data_wavelet	1✔
303
304	@property	1✔
305	def summary_dict(self) -> Dict[str, float]:	1✔
306	"""Summary dictionary of analysis metrics, including signal-to-noise ratios (SNR)."""
307	if not hasattr(self, "_summary_dict"):	1✔
308	windowed = self.highpass_fmin is not None and self.highpass_fmin > 0	1✔
309
310	self._summary_dict = dict(	1✔
311	ht=self.ht,
312	gaps=self.gaps,
313	windowed=windowed,
314	noise=self.noise,
315	**self.snr_dict,
316	)
317	self._summary_dict["lnL@true"] = f"{self.lnl(*self.waveform_parameters)}:.2e"	1✔
318
319	return self._summary_dict	1✔
320
321	@property	1✔
322	def summary(self) -> str:	1✔
323	"""Formatted summary string of analysis metrics."""
324	return "\n".join([f"{k}: {v}" for k, v in self.summary_dict.items()])	1✔
325
326	@property	1✔
327	def snr_dict(self) -> Dict[str, float]:	1✔
328	"""Calculate various SNR values based on the analysis data."""
329	if not hasattr(self, "_snr_dict"):	1✔
330	self._snr_dict = compute_snr_dict(	1✔
331	self.hf, self.psd_freqseries, self.data_frequencyseries,
332	self.hwavelet, self.psd_wavelet, self.data_wavelet,
333	self.psd_analysis, self.gaps, self.hwavelet_gapped
334	)
335	return self._snr_dict	1✔
336
337
338
339	def plot_data(self, plotfn: str=None, **kwargs):	1✔
340	"""Plot data visualizations including SNR information, time series, wavelet, and PSDs."""
341	return plot_analysis_data(self, plotfn, **kwargs)	1✔
342
343	def htemplate(self, args, *kwargs) -> Wavelet:	1✔
344	ht = self.waveform_generator(args, *kwargs, t=self.time)	1✔
345	if self.gaps is not None:	1✔
346	hwavelet = self.gaps.gap_n_transform_timeseries(	1✔
347	ht, self.Nf, self.alpha, self.highpass_fmin
348	)
349	else:
350	if self.highpass_fmin:	1!
351	ht = ht.highpass_filter(self.highpass_fmin, self.alpha)	1✔
352	hwavelet = ht.to_wavelet(Nf=self.Nf)	1✔
353	return hwavelet	1✔
354
355
356	def lnl(self, *args) -> float:	1✔
357	return compute_likelihood(	1✔
358	self.data_wavelet, self.htemplate(*args), self.psd_analysis, self.mask
359	)
360
361	def noise_lnl(self, *args) -> float:	1✔
362	return compute_likelihood(	×
363	self.data_wavelet, self.zero_wavelet, self.psd_analysis, self.mask
364	)
365
366	def ln_posterior(self, theta:List[float]) -> float:	1✔
367	if self.log_prior(np.array(theta)) == -np.inf:	×
368	return -np.inf	×
369	return self.lnl(*theta)	×
370
371	def freqdomain_lnp(self, theta:List[float]) -> float:	1✔
372	if self.log_prior(np.array(theta)) == -np.inf:	×
373	return -np.inf	×
374	return self.freqdomain_lnl(*theta)	×
375
376
377	def freqdomain_lnl(self, *args) -> float:	1✔
378	ht = self.waveform_generator(*args, t=self.time)	×
379	d = self.data_timeseries	×
380	# if self.highpass_fmin:
381	# ht = ht.highpass_filter(self.highpass_fmin, self.alpha)
382	if self.gaps is not None:	×
383	ht.data[self.gaps.gap_bools] = 0	×
384	d.data[self.gaps.gap_bools] = 0	×
385	signal_f = ht.to_frequencyseries().data	×
386	data_f = d.to_frequencyseries().data	×
387	variance_noise_f = (	×
388	self.ND * self.psd_freqseries.data / (4 * self.dt)
389	) # Calculate variance of noise, real and imaginary.
390	inn_prod = sum((abs(data_f - signal_f) ** 2) / variance_noise_f)	×
391	return -0.5 * inn_prod	×
392
393
394
395	def round_to_nearest_pow_of_2(n: int) -> int:	1✔
396	return 2 ** int(np.log2(n))	1✔
397
398
399	def get_suggested_tmax(tmax:float, dt:float=DT) -> float:	1✔
NEW 400	n = round_to_nearest_pow_of_2(tmax / dt)	×
NEW 401	return n * dt	×
402
403
404	def construct_prior(trues, ranges=None) -> ProbDistContainer:	1✔
405	ln_a, ln_f, ln_fdot = trues	1✔
406	if ranges is None:	1!
407	lna_range = [ln_a - 0.1, ln_a + 0.1]	1✔
408	lnf_range = [ln_f - 0.0001, ln_f + 0.0001]	1✔
409	lnfdot_range = [ln_fdot - 0.0001, ln_fdot + 0.0001]	1✔
410	ranges = [lna_range, lnf_range, lnfdot_range]	1✔
411	return ProbDistContainer({i: uniform_dist(*r) for i, r in enumerate(ranges)})	1✔
412
413

pywavelet / case_studies.gaps / 13365054372

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous