13128169566

Committed 04 Feb 2025 04:09AM UTC coverage: 72.733% (-0.07%) from 72.807%

Build # 13128169566

Build Type

push

github

Committed by

avivajpeyi

Commit Message

start sliurm

Coverage Stats

249 of 448 branches covered (55.58%)

Branch coverage included in aggregate %.

1 of 1 new or added line in 1 file covered. (100.0%)

12 existing lines in 1 file now uncovered.

914 of 1151 relevant lines covered (79.41%)

0.79 hits per line

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

80.0

/src/gap_study_utils/plotting/plotting.py



import arviz as az
import corner
import gif
import matplotlib.pyplot as plt
import numpy as np
from tqdm.auto import trange
from xarray.ufuncs import absolute

from ..gaps import GapType

gif.options.matplotlib["dpi"] = 100


def plot_trace(idata: az.InferenceData, axes, i=None, max_iter=None, trues=None):
    if i is not None:
        sliced_posterior = idata.posterior.isel(
            chain=slice(None), draw=slice(0, i)
        )
        idata = az.InferenceData(posterior=sliced_posterior)

    az.plot_trace(idata, axes=axes)
    for row in range(3):
        axes[row, 0].axvline(
            trues[row], c="red", linestyle="--", label="truth"
        )
        axes[row, 1].axhline(
            trues[row], c="red", linestyle="--", label="truth"
        )
        if i is not None:
            axes[row, 1].axvline(i, c="green", linestyle="--")
            axes[row, 1].set_xlim(0, max_iter)

@gif.frame
def _trace_mcmc_frame(idata, analysis_data: "AnalysisData", i, max_iter=None):
    plot_mcmc_summary(idata, analysis_data, max_iter)


def plot_mcmc_summary(idata, analysis_data: "AnalysisData", i=None, fname=None, frange=None, extra_info=""):
    if isinstance(idata, str):
        idata = az.from_netcdf(idata)

    max_iter = len(idata.sample_stats.draw)
    i = i or len(idata.sample_stats.draw) - 1
    ith_samp = idata.posterior.isel(draw=i).median(dim="chain")
    ith_samp = {
        param: float(ith_samp[param].values) for param in ith_samp.data_vars
    }
    htemplate = analysis_data.htemplate(**ith_samp)

    fig, axes = plt.subplots(4, 2, figsize=(10, 10))
    fig.suptitle(f"Iteration {i}" + f" [{extra_info}]")
    plot_trace(idata, axes, i, max_iter, trues=analysis_data.waveform_parameters)

    try:
        analysis_data.data_wavelet.plot(
            ax=axes[3, 0],
            show_colorbar=False,
            label="Whiten Data\n",
            whiten_by=analysis_data.psd_analysis.data,
            absolute=True,
            zscale="log",
        )
        htemplate.plot(
            ax=axes[3, 1],
            show_colorbar=False,
            label="ith-sample Signal\n",
            absolute=True,
            zscale="log",
        )
        if frange:
            axes[3, 0].axhline(frange[0], c="red", linestyle="--")
            axes[3, 0].axhline(frange[1], c="red", linestyle="--")
            axes[3, 1].axhline(frange[0], c="red", linestyle="--")
            axes[3, 1].axhline(frange[1], c="red", linestyle="--")
    except Exception as e:
        print(e)


    if fname:
        plt.tight_layout()
        plt.savefig(fname)


def make_mcmc_trace_gif(
    idata_fname, analysis_data, n_frames=20, fname="mcmc.gif"
):
    idata = az.from_netcdf(idata_fname)
    trace_frames = []
    N = len(idata.sample_stats.draw)
    trace_frames.append(_trace_mcmc_frame(idata, analysis_data, 1, N))
    for i in trange(int(N * 0.1), N, int(N / n_frames)):
        trace_frames.append(_trace_mcmc_frame(idata, analysis_data, i, N))
    gif.save(trace_frames, fname, duration=100)




def plot_analysis_data(analysis_data:"AnalysisData", plotfn:str, compact=True, figsize=(5, 8)):
    fig, ax = plt.subplots(6, 1, figsize=figsize)
    ax[0].axis("off")
    ax[0].text(
        0.0, 0.5, analysis_data.summary, fontsize=6, verticalalignment="center"
    )

    plot_analysis_fseries(analysis_data, ax[1])
    plot_analysis_tseries(analysis_data, ax[2])
    plot_analysis_wdm(analysis_data, ax[3:])
    if compact:
        plt.subplots_adjust(hspace=0)
    else:
        plt.tight_layout()
    if plotfn:
        plt.savefig(plotfn, bbox_inches="tight")


def plot_analysis_fseries(data:"AnalysisData", ax):
    data.hf.plot_periodogram(ax=ax, color="C0", alpha=1, lw=1)
    data.psd_freqseries.plot(ax=ax, color="k", alpha=1, lw=1)
    if data.highpass_fmin:
        ax.set_xlim(left=data.highpass_fmin)
    ax.set_xlim(right=data.freq[-1])
    ax.tick_params(
        axis="x",
        direction="in",
        labelbottom=False,
        top=True,
        labeltop=True,
    )
    if data.gaps and data.gaps.type == GapType.STITCH:
            for i, chunkf_i in enumerate(data.chunked_hf):
                chunkf_i.plot_periodogram(
                    ax=ax, color=f"C{i + 1}", alpha=0.5
                )
    ax.set_ylabel("PSD [Hz$^{-1}$]")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_yscale("log")
    ax.set_xscale("log")


def plot_analysis_tseries(data:"AnalysisData", ax):
    data.ht.plot(ax=ax)
    if data.gaps:
        for i, chunk_i in enumerate(data.chunked_ht):
            chunk_i.plot(ax=ax, color=f"C{i + 1}", alpha=0.5)
    ax.set_ylabel("Strain")
    ax.set_xlabel("Time [s]")
    ax.set_xlim(0, data.tmax)

def plot_analysis_wdm(data:"AnalysisData", axes, whiten=True):
    kwgs_psd = dict(show_colorbar=False, absolute=True, zscale="log")
    kwgs = dict(show_colorbar=False)
    data_label, model_label = "Data", "Model"
    if whiten:
        kwgs = dict(whiten_by=data.psd_wavelet.data, absolute=True, **kwgs)
        data_label = "Whitened "  + data_label
        model_label = "Whitened " + model_label

    data_w = data.data_wavelet
    hwavelet = data.hwavelet
    psd_w = data.psd_wavelet
    if data.mask:
        data_w = data_w * data.mask
        hwavelet = hwavelet * data.mask
        psd_w = psd_w * data.mask


    data_w.plot(ax=axes[0], label=data_label, **kwgs)
    hwavelet.plot(ax=axes[1], label=model_label, **kwgs)
    psd_w.plot(ax=axes[2], label="PSD", **kwgs_psd)


    for ax in axes:
        for f in data.frange:
            ax.axhline(f, color="r", linestyle="--")
        if data.gaps:
            data.gaps.plot(ax=ax)

1
2
3	import arviz as az	1✔
4	import corner	1✔
5	import gif	1✔
6	import matplotlib.pyplot as plt	1✔
7	import numpy as np	1✔
8	from tqdm.auto import trange	1✔
9	from xarray.ufuncs import absolute	1✔
10
11	from ..gaps import GapType	1✔
12
13	gif.options.matplotlib["dpi"] = 100	1✔
14
15
16	def plot_trace(idata: az.InferenceData, axes, i=None, max_iter=None, trues=None):	1✔
17	if i is not None:	1!
18	sliced_posterior = idata.posterior.isel(	1✔
19	chain=slice(None), draw=slice(0, i)
20	)
21	idata = az.InferenceData(posterior=sliced_posterior)	1✔
22
23	az.plot_trace(idata, axes=axes)	1✔
24	for row in range(3):	1✔
25	axes[row, 0].axvline(	1✔
26	trues[row], c="red", linestyle="--", label="truth"
27	)
28	axes[row, 1].axhline(	1✔
29	trues[row], c="red", linestyle="--", label="truth"
30	)
31	if i is not None:	1!
32	axes[row, 1].axvline(i, c="green", linestyle="--")	1✔
33	axes[row, 1].set_xlim(0, max_iter)	1✔
34
35	@gif.frame	1✔
36	def _trace_mcmc_frame(idata, analysis_data: "AnalysisData", i, max_iter=None):	1✔
37	plot_mcmc_summary(idata, analysis_data, max_iter)	×
38
39
40	def plot_mcmc_summary(idata, analysis_data: "AnalysisData", i=None, fname=None, frange=None, extra_info=""):	1✔
41	if isinstance(idata, str):	1!
42	idata = az.from_netcdf(idata)	1✔
43
44	max_iter = len(idata.sample_stats.draw)	1✔
45	i = i or len(idata.sample_stats.draw) - 1	1✔
46	ith_samp = idata.posterior.isel(draw=i).median(dim="chain")	1✔
47	ith_samp = {	1✔
48	param: float(ith_samp[param].values) for param in ith_samp.data_vars
49	}
50	htemplate = analysis_data.htemplate(**ith_samp)	1✔
51
52	fig, axes = plt.subplots(4, 2, figsize=(10, 10))	1✔
53	fig.suptitle(f"Iteration {i}" + f" [{extra_info}]")	1✔
54	plot_trace(idata, axes, i, max_iter, trues=analysis_data.waveform_parameters)	1✔
55
56	try:	1✔
57	analysis_data.data_wavelet.plot(	1✔
58	ax=axes[3, 0],
59	show_colorbar=False,
60	label="Whiten Data\n",
61	whiten_by=analysis_data.psd_analysis.data,
62	absolute=True,
63	zscale="log",
64	)
65	htemplate.plot(	1✔
66	ax=axes[3, 1],
67	show_colorbar=False,
68	label="ith-sample Signal\n",
69	absolute=True,
70	zscale="log",
71	)
72	if frange:	1!
73	axes[3, 0].axhline(frange[0], c="red", linestyle="--")	×
74	axes[3, 0].axhline(frange[1], c="red", linestyle="--")	×
UNCOV 75	axes[3, 1].axhline(frange[0], c="red", linestyle="--")	×
UNCOV 76	axes[3, 1].axhline(frange[1], c="red", linestyle="--")	×
UNCOV 77	except Exception as e:	×
UNCOV 78	print(e)	×
79
80
81	if fname:	1!
82	plt.tight_layout()	1✔
83	plt.savefig(fname)	1✔
84
85
86	def make_mcmc_trace_gif(	1✔
87	idata_fname, analysis_data, n_frames=20, fname="mcmc.gif"
88	):
89	idata = az.from_netcdf(idata_fname)	×
UNCOV 90	trace_frames = []	×
UNCOV 91	N = len(idata.sample_stats.draw)	×
UNCOV 92	trace_frames.append(_trace_mcmc_frame(idata, analysis_data, 1, N))	×
UNCOV 93	for i in trange(int(N * 0.1), N, int(N / n_frames)):	×
UNCOV 94	trace_frames.append(_trace_mcmc_frame(idata, analysis_data, i, N))	×
UNCOV 95	gif.save(trace_frames, fname, duration=100)	×
96
97
98
99
100	def plot_analysis_data(analysis_data:"AnalysisData", plotfn:str, compact=True, figsize=(5, 8)):	1✔
101	fig, ax = plt.subplots(6, 1, figsize=figsize)	1✔
102	ax[0].axis("off")	1✔
103	ax[0].text(	1✔
104	0.0, 0.5, analysis_data.summary, fontsize=6, verticalalignment="center"
105	)
106
107	plot_analysis_fseries(analysis_data, ax[1])	1✔
108	plot_analysis_tseries(analysis_data, ax[2])	1✔
109	plot_analysis_wdm(analysis_data, ax[3:])	1✔
110	if compact:	1!
111	plt.subplots_adjust(hspace=0)	1✔
112	else:
UNCOV 113	plt.tight_layout()	×
114	if plotfn:	1!
115	plt.savefig(plotfn, bbox_inches="tight")	1✔
116
117
118	def plot_analysis_fseries(data:"AnalysisData", ax):	1✔
119	data.hf.plot_periodogram(ax=ax, color="C0", alpha=1, lw=1)	1✔
120	data.psd_freqseries.plot(ax=ax, color="k", alpha=1, lw=1)	1✔
121	if data.highpass_fmin:	1!
UNCOV 122	ax.set_xlim(left=data.highpass_fmin)	×
123	ax.set_xlim(right=data.freq[-1])	1✔
124	ax.tick_params(	1✔
125	axis="x",
126	direction="in",
127	labelbottom=False,
128	top=True,
129	labeltop=True,
130	)
131	if data.gaps and data.gaps.type == GapType.STITCH:	1✔
132	for i, chunkf_i in enumerate(data.chunked_hf):	1✔
133	chunkf_i.plot_periodogram(	1✔
134	ax=ax, color=f"C{i + 1}", alpha=0.5
135	)
136	ax.set_ylabel("PSD [Hz$^{-1}$]")	1✔
137	ax.set_xlabel("Frequency [Hz]")	1✔
138	ax.set_yscale("log")	1✔
139	ax.set_xscale("log")	1✔
140
141
142	def plot_analysis_tseries(data:"AnalysisData", ax):	1✔
143	data.ht.plot(ax=ax)	1✔
144	if data.gaps:	1!
145	for i, chunk_i in enumerate(data.chunked_ht):	1✔
146	chunk_i.plot(ax=ax, color=f"C{i + 1}", alpha=0.5)	1✔
147	ax.set_ylabel("Strain")	1✔
148	ax.set_xlabel("Time [s]")	1✔
149	ax.set_xlim(0, data.tmax)	1✔
150
151	def plot_analysis_wdm(data:"AnalysisData", axes, whiten=True):	1✔
152	kwgs_psd = dict(show_colorbar=False, absolute=True, zscale="log")	1✔
153	kwgs = dict(show_colorbar=False)	1✔
154	data_label, model_label = "Data", "Model"	1✔
155	if whiten:	1!
156	kwgs = dict(whiten_by=data.psd_wavelet.data, absolute=True, **kwgs)	1✔
157	data_label = "Whitened " + data_label	1✔
158	model_label = "Whitened " + model_label	1✔
159
160	data_w = data.data_wavelet	1✔
161	hwavelet = data.hwavelet	1✔
162	psd_w = data.psd_wavelet	1✔
163	if data.mask:	1!
164	data_w = data_w * data.mask	1✔
165	hwavelet = hwavelet * data.mask	1✔
166	psd_w = psd_w * data.mask	1✔
167
168
169	data_w.plot(ax=axes[0], label=data_label, **kwgs)	1✔
170	hwavelet.plot(ax=axes[1], label=model_label, **kwgs)	1✔
171	psd_w.plot(ax=axes[2], label="PSD", **kwgs_psd)	1✔
172
173
174	for ax in axes:	1✔
175	for f in data.frange:	1✔
176	ax.axhline(f, color="r", linestyle="--")	1✔
177	if data.gaps:	1!
178	data.gaps.plot(ax=ax)	1✔

pywavelet / case_studies.gaps / 13128169566

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