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

91.89

/src/flowMC/Sampler.py

import jax.numpy as jnp
from jaxtyping import Array, Float, PRNGKeyArray

from flowMC.strategy.base import Strategy
from flowMC.resource.base import Resource
from flowMC.resource_strategy_bundles import ResourceStrategyBundle


class Sampler:
    """Top level API that the users primarily interact with.

    Args:
        n_dim (int): Dimension of the parameter space.
        n_chains (int): Number of chains to sample.
        rng_key (PRNGKeyArray): Jax PRNGKey.
        logpdf (Callable[[Float[Array, "n_dim"], dict], Float): Log probability function.
        resources (dict[str, Resource]): Resources to be used by the sampler.
        strategies (list[Strategy]): List of strategies to be used by the sampler.
        verbose (bool): Whether to print out progress. Defaults to False.
        logging (bool): Whether to log the progress. Defaults to True.
        outdir (str): Directory to save the logs. Defaults to "./outdir/".
    """

    # Essential parameters
    n_dim: int
    n_chains: int
    rng_key: PRNGKeyArray
    resources: dict[str, Resource]
    strategies: list[Strategy]

    # Logging hyperparameters
    verbose: bool = False
    logging: bool = True
    outdir: str = "./outdir/"

    def __init__(
        self,
        n_dim: int,
        n_chains: int,
        rng_key: PRNGKeyArray,
        resources: None | dict[str, Resource] = None,
        strategies: None | list[Strategy] = None,
        resource_strategy_bundles: None | ResourceStrategyBundle = None,
        **kwargs,
    ):
        # Copying input into the model

        self.n_dim = n_dim
        self.n_chains = n_chains
        self.rng_key = rng_key

        if resources is not None and strategies is not None:
            print(
                "Resources and strategies provided. Ignoring resource strategy bundles."
            )
            self.resources = resources
            self.strategies = strategies
        else:
            print(
                "Resources or strategies not provided. Using resource strategy bundles."
            )
            assert (
                resource_strategy_bundles is not None
            ), "Resource strategy bundles must be provided if resources and strategies are not."
            self.resources = resource_strategy_bundles.resources
            self.strategies = resource_strategy_bundles.strategies

        # Set and override any given hyperparameters
        class_keys = list(self.__class__.__dict__.keys())
        for key, value in kwargs.items():
            if key in class_keys:
                if not key.startswith("__"):
                    setattr(self, key, value)

    def sample(self, initial_position: Float[Array, "n_chains n_dim"], data: dict):
        """Sample from the posterior using the local sampler.

        Args:
            initial_position (Device Array): Initial position.
            data (dict): Data to be used by the likelihood functions
        """

        initial_position = jnp.atleast_2d(initial_position)  # type: ignore
        rng_key = self.rng_key
        last_step = initial_position
        for strategy in self.strategies:
            (
                rng_key,
                self.resources,
                last_step,
            ) = strategy(rng_key, self.resources, last_step, data)

    # TODO: Implement quick access and summary functions that operates on buffer

1	import jax.numpy as jnp	2✔
2	from jaxtyping import Array, Float, PRNGKeyArray	2✔
3
4	from flowMC.strategy.base import Strategy	2✔
5	from flowMC.resource.base import Resource	2✔
6	from flowMC.resource_strategy_bundles import ResourceStrategyBundle	2✔
7
8
9	class Sampler:	2✔
10	"""Top level API that the users primarily interact with.
11
12	Args:
13	n_dim (int): Dimension of the parameter space.
14	n_chains (int): Number of chains to sample.
15	rng_key (PRNGKeyArray): Jax PRNGKey.
16	logpdf (Callable[[Float[Array, "n_dim"], dict], Float): Log probability function.
17	resources (dict[str, Resource]): Resources to be used by the sampler.
18	strategies (list[Strategy]): List of strategies to be used by the sampler.
19	verbose (bool): Whether to print out progress. Defaults to False.
20	logging (bool): Whether to log the progress. Defaults to True.
21	outdir (str): Directory to save the logs. Defaults to "./outdir/".
22	"""
23
24	# Essential parameters
25	n_dim: int	2✔
26	n_chains: int	2✔
27	rng_key: PRNGKeyArray	2✔
28	resources: dict[str, Resource]	2✔
29	strategies: list[Strategy]	2✔
30
31	# Logging hyperparameters
32	verbose: bool = False	2✔
33	logging: bool = True	2✔
34	outdir: str = "./outdir/"	2✔
35
36	def __init__(	2✔
37	self,
38	n_dim: int,
39	n_chains: int,
40	rng_key: PRNGKeyArray,
41	resources: None \| dict[str, Resource] = None,
42	strategies: None \| list[Strategy] = None,
43	resource_strategy_bundles: None \| ResourceStrategyBundle = None,
44	**kwargs,
45	):
46	# Copying input into the model
47
48	self.n_dim = n_dim	2✔
49	self.n_chains = n_chains	2✔
50	self.rng_key = rng_key	2✔
51
52	if resources is not None and strategies is not None:	2✔
53	print(	2✔
54	"Resources and strategies provided. Ignoring resource strategy bundles."
55	)
56	self.resources = resources	2✔
57	self.strategies = strategies	2✔
58	else:
59	print(	2✔
60	"Resources or strategies not provided. Using resource strategy bundles."
61	)
62	assert (	2✔
63	resource_strategy_bundles is not None
64	), "Resource strategy bundles must be provided if resources and strategies are not."
65	self.resources = resource_strategy_bundles.resources	2✔
66	self.strategies = resource_strategy_bundles.strategies	2✔
67
68	# Set and override any given hyperparameters
69	class_keys = list(self.__class__.__dict__.keys())	2✔
70	for key, value in kwargs.items():	2✔
UNCOV 71	if key in class_keys:	×
72	if not key.startswith("__"):	×
73	setattr(self, key, value)	×
74
75	def sample(self, initial_position: Float[Array, "n_chains n_dim"], data: dict):	2✔
76	"""Sample from the posterior using the local sampler.
77
78	Args:
79	initial_position (Device Array): Initial position.
80	data (dict): Data to be used by the likelihood functions
81	"""
82
83	initial_position = jnp.atleast_2d(initial_position) # type: ignore	2✔
84	rng_key = self.rng_key	2✔
85	last_step = initial_position	2✔
86	for strategy in self.strategies:	2✔
87	(	2✔
88	rng_key,
89	self.resources,
90	last_step,
91	) = strategy(rng_key, self.resources, last_step, data)
92
93	# TODO: Implement quick access and summary functions that operates on buffer

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