13655281414

Committed 04 Mar 2025 01:55PM UTC coverage: 67.835%. Remained the same

Build # 13655281414

Build Type

push

github

Committed by

kazewong

Commit Message

format doc strings

Run Details

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

162 existing lines in 15 files now uncovered.

987 of 1455 relevant lines covered (67.84%)

1.36 hits per line

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

77.78

/src/flowMC/resource/local_kernel/base.py

from abc import abstractmethod
from typing import Callable

import jax
from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
from flowMC.resource.base import Resource


@jax.tree_util.register_pytree_node_class
class ProposalBase(Resource):
    def __init__(
        self,
        **kwargs,
    ):
        """Initialize the sampler class."""
        self.kwargs = kwargs

    @abstractmethod
    def kernel(
        self,
        rng_key: PRNGKeyArray,
        log_pdf: Callable[[Float[Array, " n_dim"], PyTree], Float[Array, "1"]],
        position: Float[Array, "nstep  n_dim"],
        log_prob: Float[Array, "nstep 1"],
        data: PyTree,
    ) -> tuple[
        Float[Array, "nstep  n_dim"], Float[Array, "nstep 1"], Int[Array, "n_step 1"]
    ]:
        """Kernel for one step in the proposal cycle."""

    def tree_flatten(self):
        children = ()

        aux_data = {"kwargs": self.kwargs}
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children, **aux_data)

1	from abc import abstractmethod	2✔
2	from typing import Callable	2✔
3
4	import jax	2✔
5	from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree	2✔
6	from flowMC.resource.base import Resource	2✔
7
8
9	@jax.tree_util.register_pytree_node_class	2✔
10	class ProposalBase(Resource):	2✔
11	def __init__(	2✔
12	self,
13	**kwargs,
14	):
15	"""Initialize the sampler class."""
16	self.kwargs = kwargs	2✔
17
18	@abstractmethod	2✔
19	def kernel(	2✔
20	self,
21	rng_key: PRNGKeyArray,
22	log_pdf: Callable[[Float[Array, " n_dim"], PyTree], Float[Array, "1"]],
23	position: Float[Array, "nstep n_dim"],
24	log_prob: Float[Array, "nstep 1"],
25	data: PyTree,
26	) -> tuple[
27	Float[Array, "nstep n_dim"], Float[Array, "nstep 1"], Int[Array, "n_step 1"]
28	]:
29	"""Kernel for one step in the proposal cycle."""
30
31	def tree_flatten(self):	2✔
UNCOV 32	children = ()	×
33
UNCOV 34	aux_data = {"kwargs": self.kwargs}	×
UNCOV 35	return (children, aux_data)	×
36
37	@classmethod	2✔
38	def tree_unflatten(cls, aux_data, children):	2✔
39	return cls(children, *aux_data)	×

kazewong / flowMC / 13655281414

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