18303071607

Committed 07 Oct 2025 05:32AM UTC coverage: 80.332% (-0.6%) from 80.952%

Build # 18303071607

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #10 from nz-gravity/save_vi_diagnostics_before_sampling

save VI plots at the start

Run Details

576 of 694 branches covered (83.0%)

Branch coverage included in aggregate %.

63 of 80 new or added lines in 7 files covered. (78.75%)

498 existing lines in 11 files now uncovered.

3925 of 4909 relevant lines covered (79.96%)

1.6 hits per line

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

60.29

/src/log_psplines/plotting/utils.py

import dataclasses

import jax.numpy as jnp
import numpy as np

from ..datatypes import Periodogram
from ..psplines import LogPSplines
from .base import compute_confidence_intervals, subsample_weights

__all__ = ["unpack_data"]


@dataclasses.dataclass
class PlottingData:
    freqs: np.ndarray = None
    pdgrm: np.ndarray = None
    model: np.ndarray = None
    ci: np.ndarray = None

    @property
    def n(self):
        if self.freqs is not None:
            return len(self.freqs)
        elif self.pdgrm is not None:
            return len(self.pdgrm)
        elif self.model is not None:
            return len(self.model)
        else:
            raise ValueError("No data to get length from.")


def unpack_data(
    pdgrm: Periodogram = None,
    spline_model: LogPSplines = None,
    weights=None,
    yscalar=1.0,
    use_uniform_ci=True,
    use_parametric_model=True,
    freqs=None,
    posterior_psd=None,
    model_ci=None,
):
    plt_dat = PlottingData()
    if pdgrm is not None:
        plt_dat.pdgrm = np.array(pdgrm.power, dtype=np.float64) * yscalar
        plt_dat.freqs = np.array(pdgrm.freqs)

    if plt_dat.freqs is None and freqs is None:
        plt_dat.freqs = np.linspace(0, 1, plt_dat.n)
    elif freqs is not None:
        plt_dat.freqs = freqs

    if model_ci is not None:
        plt_dat.ci = model_ci
        plt_dat.model = model_ci[1]

    if posterior_psd is not None:
        ci = np.percentile(posterior_psd, q=jnp.array([16, 50, 84]), axis=0)
        plt_dat.ci = ci
        plt_dat.model = ci[1]

    if plt_dat.model is None and spline_model is not None:

        if weights is None:
            # just use the initial weights/0 weights
            ln_spline = spline_model(use_parametric_model=use_parametric_model)

        elif weights.ndim == 1:
            # only one set of weights -- no CI possible
            ln_spline = spline_model(weights, use_parametric_model)

        else:  # weights.ndim == 2
            # multiple sets of weights -- CI possible

            if weights.shape[0] > 500:
                # subsample to speed up
                idx = np.random.choice(
                    weights.shape[0], size=500, replace=False
                )
                weights = weights[idx]

            ln_splines = jnp.array(
                [spline_model(w, use_parametric_model) for w in weights]
            )

            if use_uniform_ci:
                ln_ci = compute_confidence_intervals(
                    ln_splines, method="uniform"
                )
            else:  # percentile
                ln_ci = compute_confidence_intervals(
                    ln_splines, method="percentile"
                )
            ln_ci = jnp.array(ln_ci)
            plt_dat.ci = np.exp(ln_ci, dtype=np.float64) * yscalar
            ln_spline = ln_ci[1]
        plt_dat.model = np.exp(ln_spline, dtype=np.float64) * yscalar

    return plt_dat

1	import dataclasses	2✔
2
3	import jax.numpy as jnp	2✔
4	import numpy as np	2✔
5
6	from ..datatypes import Periodogram	2✔
7	from ..psplines import LogPSplines	2✔
8	from .base import compute_confidence_intervals, subsample_weights	2✔
9
10	__all__ = ["unpack_data"]	2✔
11
12
13	@dataclasses.dataclass	2✔
14	class PlottingData:	2✔
15	freqs: np.ndarray = None	2✔
16	pdgrm: np.ndarray = None	2✔
17	model: np.ndarray = None	2✔
18	ci: np.ndarray = None	2✔
19
20	@property	2✔
21	def n(self):	2✔
22	if self.freqs is not None:	×
23	return len(self.freqs)	×
24	elif self.pdgrm is not None:	×
25	return len(self.pdgrm)	×
26	elif self.model is not None:	×
UNCOV 27	return len(self.model)	×
28	else:
UNCOV 29	raise ValueError("No data to get length from.")	×
30
31
32	def unpack_data(	2✔
33	pdgrm: Periodogram = None,
34	spline_model: LogPSplines = None,
35	weights=None,
36	yscalar=1.0,
37	use_uniform_ci=True,
38	use_parametric_model=True,
39	freqs=None,
40	posterior_psd=None,
41	model_ci=None,
42	):
43	plt_dat = PlottingData()	2✔
44	if pdgrm is not None:	2✔
45	plt_dat.pdgrm = np.array(pdgrm.power, dtype=np.float64) * yscalar	2✔
46	plt_dat.freqs = np.array(pdgrm.freqs)	2✔
47
48	if plt_dat.freqs is None and freqs is None:	2✔
49	plt_dat.freqs = np.linspace(0, 1, plt_dat.n)	×
50	elif freqs is not None:	2✔
UNCOV 51	plt_dat.freqs = freqs	×
52
53	if model_ci is not None:	2✔
54	plt_dat.ci = model_ci	2✔
55	plt_dat.model = model_ci[1]	2✔
56
57	if posterior_psd is not None:	2✔
58	ci = np.percentile(posterior_psd, q=jnp.array([16, 50, 84]), axis=0)	2✔
59	plt_dat.ci = ci	2✔
60	plt_dat.model = ci[1]	2✔
61
62	if plt_dat.model is None and spline_model is not None:	2✔
63
64	if weights is None:	2✔
65	# just use the initial weights/0 weights
66	ln_spline = spline_model(use_parametric_model=use_parametric_model)	2✔
67
UNCOV 68	elif weights.ndim == 1:	×
69	# only one set of weights -- no CI possible
UNCOV 70	ln_spline = spline_model(weights, use_parametric_model)	×
71
72	else: # weights.ndim == 2
73	# multiple sets of weights -- CI possible
74
UNCOV 75	if weights.shape[0] > 500:	×
76	# subsample to speed up
UNCOV 77	idx = np.random.choice(	×
78	weights.shape[0], size=500, replace=False
79	)
80	weights = weights[idx]	×
81
UNCOV 82	ln_splines = jnp.array(	×
83	[spline_model(w, use_parametric_model) for w in weights]
84	)
85
86	if use_uniform_ci:	×
87	ln_ci = compute_confidence_intervals(	×
88	ln_splines, method="uniform"
89	)
90	else: # percentile
UNCOV 91	ln_ci = compute_confidence_intervals(	×
92	ln_splines, method="percentile"
93	)
UNCOV 94	ln_ci = jnp.array(ln_ci)	×
UNCOV 95	plt_dat.ci = np.exp(ln_ci, dtype=np.float64) * yscalar	×
UNCOV 96	ln_spline = ln_ci[1]	×
97	plt_dat.model = np.exp(ln_spline, dtype=np.float64) * yscalar	2✔
98
99	return plt_dat	2✔

nz-gravity / LogPSplinePSD / 18303071607

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