compiled-knowledge 4.0.0a5__cp313-cp313-macosx_10_13_universal2.whl

ck/sampling/sampler_support.py

@@ -0,0 +1,232 @@
+from dataclasses import dataclass
+from itertools import count
+from typing import Callable, Sequence, Optional, Set, Tuple, Dict, Collection
+
+from ck.pgm import Instance, RandomVariable, Indicator
+from ck.pgm_circuit.program_with_slotmap import ProgramWithSlotmap
+from ck.pgm_circuit.slot_map import SlotMap
+from ck.probability.probability_space import Condition, check_condition
+from ck.utils.map_set import MapSet
+from ck.utils.np_extras import NDArrayStates, NDArrayNumeric
+from ck.utils.random_extras import Random
+
+# Type of a yield function. Support for a sampler.
+# A yield function may be used to implement a sampler's iterator, thus
+# it provides an Instance or single state index.
+YieldF = Callable[[NDArrayStates], Instance | int]
+
+
+@dataclass
+class SampleRV:
+    """
+    Support for a sampler.
+    A SampleRV structure keeps track of information for one sampled random variable.
+    """
+    index: int  # index into the sequence of sample rvs.
+    rv: RandomVariable  # the random variable being sampled.
+    slots: Sequence[int]  # program input slots for indicators of the random variable (co-indexed with rv.states).
+    copy_index: Optional[int]  # for Markov chains, which previous sample rv should be copied?
+
+
+@dataclass
+class SamplerInfo:
+    """
+    Support for a sampler.
+    A SamplerInfo structure keeps track of standard information when a sampler uses a Program.
+    """
+    sample_rvs: Sequence[SampleRV]
+    condition: Sequence[Indicator]
+    yield_f: YieldF
+    slots_0: Set[int]
+    slots_1: Set[int]
+
+    @property
+    def rvs(self) -> Tuple[RandomVariable, ...]:
+        """
+        Extract the RandomVariable objects from `self.sample_rvs`.
+        """
+        return tuple(sample_rv.rv for sample_rv in self.sample_rvs)
+
+
+def get_sampler_info(
+        program_with_slotmap: ProgramWithSlotmap,
+        rvs: Optional[RandomVariable | Sequence[RandomVariable]],
+        condition: Condition,
+        chain_pairs: Sequence[Tuple[RandomVariable, RandomVariable]] = (),
+        initial_chain_condition: Condition = (),
+) -> SamplerInfo:
+    """
+    Helper for samplers.
+
+    Determines:
+    (1) the slots for sampling rvs,
+    (2) Markov chaining rvs,
+    (3) the function to use for yielding an Instance or state index.
+
+    If parameter `rvs` is a RandomVariable, then the yield function will
+    provide a state index. If parameter `rvs` is a Sequence, then the
+    yield function will provide an Instance.
+
+    Args:
+        program_with_slotmap: the program and slotmap being referenced.
+        rvs: the random variables to sample. It may be either a sequence of
+            random variables, or a single random variable.
+        condition: is a collection of zero or more conditioning indicators.
+        chain_pairs: is a collection of pairs of random variables, each random variable
+            must be in the given rvs. Given a pair (from_rv, to_rv) the state of from_rv is used
+            as a condition for to_rv prior to generating a sample.
+        initial_chain_condition: are condition indicators (just like condition)
+            for the initialisation of the 'to_rv' random variables mentioned in chain_pairs.
+
+    Raises:
+        ValueError: if preconditions of `program_with_slotmap` are incompatible with the given condition.
+
+    Returns:
+        a SamplerInfo structure.
+    """
+    if rvs is None:
+        rvs = program_with_slotmap.rvs
+    if isinstance(rvs, RandomVariable):
+        # a single rv
+        rvs = (rvs,)
+        yield_f = lambda x: x.item()
+    else:
+        # a sequence of rvs
+        rvs = tuple(rvs)
+        yield_f = lambda x: x.tolist()
+
+    # Group condition indicators by `rv_idx`.
+    conditioned_rvs: MapSet[int, Indicator] = MapSet()
+    for ind in check_condition(condition):
+        conditioned_rvs.add(ind.rv_idx, ind)
+    del condition
+
+    # Group precondition indicators by `rv_idx`.
+    preconditioned_rvs: MapSet[int, Indicator] = MapSet()
+    for ind in program_with_slotmap.precondition:
+        preconditioned_rvs.add(ind.rv_idx, ind)
+
+    # Rationalise conditioned_rvs with preconditioned_rvs
+    rv_idx: int
+    precondition_set: Set[Indicator]
+    for rv_idx, precondition_set in preconditioned_rvs.items():
+        condition_set = conditioned_rvs.get(rv_idx)
+        if condition_set is None:
+            # A preconditioned rv was not mentioned in the explicit conditions
+            conditioned_rvs.add_all(rv_idx, precondition_set)
+        else:
+            # A preconditioned rv was also mentioned in the explicit conditions
+            condition_set.intersection_update(precondition_set)
+            if len(condition_set) == 0:
+                rv_index: Dict[int, RandomVariable] = {rv.idx: rv for rv in rvs}
+                rv: RandomVariable = rv_index[rv_idx]
+                raise ValueError(f'conditions on rv {rv} are disjoint from preconditions')
+    del preconditioned_rvs
+
+    # Group initial chain indicators by `rv_idx`.
+    initial_chain_condition: Sequence[Indicator] = check_condition(initial_chain_condition)
+    initial_chain_conditioned_rvs: MapSet[int, Indicator] = MapSet()
+    for ind in initial_chain_condition:
+        initial_chain_conditioned_rvs.add(ind.rv_idx, ind)
+
+    # Check sample rvs are valid and without duplicates.
+    rvs_set: Set[RandomVariable] = set(rvs)
+    if not rvs_set.issubset(program_with_slotmap.rvs):
+        raise ValueError('sample random variables not available')
+    if len(rvs) != len(rvs_set):
+        raise ValueError('duplicate sample random variables requested')
+
+    # Check chain_pairs rvs are being sampled
+    if not rvs_set.issuperset(pair[0] for pair in chain_pairs):
+        raise ValueError('a random variable appears in chain_pairs but not in sample rvs')
+    if not rvs_set.issuperset(pair[1] for pair in chain_pairs):
+        raise ValueError('a random variable appears in chain_pairs but not in sample rvs')
+
+    # Check chain_pairs source and destination rvs are disjoint
+    if not {pair[0] for pair in chain_pairs}.isdisjoint(pair[1] for pair in chain_pairs):
+        raise ValueError('chain_pairs sources and destinations are not disjoint')
+
+    # Check no chain_pairs destination rv is a conditioned rv
+    if any(pair[1].idx in conditioned_rvs.keys() for pair in chain_pairs):
+        raise ValueError('a chain_pairs destination is conditioned')
+
+    # Check chain initial conditions relate to chain_pairs destination rvs
+    chain_dest_rv_idxs: Set[int] = {pair[1].idx for pair in chain_pairs}
+    if not all(rv_idx in chain_dest_rv_idxs for rv_idx in initial_chain_conditioned_rvs.keys()):
+        raise ValueError('a chain initial condition is not a chain destination rv')
+
+    # Convert chain_pairs for registering with `sample_rvs`.
+    # rv_idx maps RandomVariable id to a position it exists in rvs (doesn't matter if rv is duplicated in rvs)
+    # copy_idx RandomVariable id to a position in rvs that it can be copied from for Markov chaining.
+    rv_idx: Dict[int, int] = {id(rv): i for i, rv in enumerate(rvs)}
+    copy_idx: Dict[int, int] = {id(rv): rv_idx[id(prev_rv)] for prev_rv, rv in chain_pairs}
+
+    # Get rv state slots, rvs_slots is co-indexed with rvs
+    slot_map: SlotMap = program_with_slotmap.slot_map
+    rvs_slots = tuple(tuple(slot_map[ind] for ind in rv) for rv in rvs)
+
+    sample_rvs: Sequence[SampleRV] = tuple(
+        SampleRV(idx, rv, rv_slots, copy_idx.get(id(rv)))
+        for idx, rv, rv_slots in zip(count(), rvs, rvs_slots)
+    )
+
+    # Process the condition to get zero and one slots
+    slots_0: Set[int] = set()
+    slots_1: Set[int] = set()
+    for rv in program_with_slotmap.rvs:
+        conditioning: Optional[Set[Indicator]] = conditioned_rvs.get(rv.idx)
+        if conditioning is not None:
+            slots_1.update(slot_map[ind] for ind in conditioning)
+            slots_0.update(slot_map[ind] for ind in rv if ind not in conditioning)
+            continue
+
+        conditioning: Optional[Set[Indicator]] = initial_chain_conditioned_rvs.get(rv.idx)
+        if conditioning is not None:
+            slots_1.update(slot_map[ind] for ind in conditioning)
+            slots_0.update(slot_map[ind] for ind in rv if ind not in conditioning)
+            continue
+
+        # default
+        slots_1.update(slot_map[ind] for ind in rv)
+
+    return SamplerInfo(
+        sample_rvs=sample_rvs,
+        condition=tuple(ind for condition_set in conditioned_rvs.values() for ind in condition_set),
+        yield_f=yield_f,
+        slots_0=slots_0,
+        slots_1=slots_1,
+    )
+
+
+def uniform_random_sample(
+        sample_rvs: Sequence[SampleRV],
+        slots_0: Collection[int],
+        slots_1: Collection[int],
+        slots: NDArrayNumeric,
+        state: NDArrayStates,
+        rand: Random,
+):
+    """
+    Helper for samplers.
+
+    Sets the states to a random instance and configures slots to match.
+    States are drawn from a uniform distribution, drawn using random.randrange.
+    """
+
+    # Set up the input slots to respect conditioning
+    for slot in slots_0:
+        slots[slot] = 0
+    for slot in slots_1:
+        slots[slot] = 1
+
+    for sample_rv in sample_rvs:
+        candidates = []
+        for slot_state, slot in enumerate(sample_rv.slots):
+            if slots[slot] == 1:
+                slots[slot] = 0
+                candidates.append((slot_state, slot))
+
+        # Pick a random state for sample_rv
+        slot_state, slot = candidates[rand.randrange(0, len(candidates))]
+        state[sample_rv.index] = slot_state
+        slots[slot] = 1

ck/sampling/uniform_sampler.py

@@ -0,0 +1,66 @@
+from typing import Set, List, Iterator, Optional, Sequence
+import random
+
+import numpy as np
+
+from ck.pgm import Instance, RandomVariable, Indicator
+from ck.probability.probability_space import dtype_for_state_indexes, Condition, check_condition
+from .sampler import Sampler
+from .sampler_support import YieldF
+from ck.utils.map_set import MapSet
+from ck.utils.np_extras import DType
+from ck.utils.random_extras import Random
+
+
+class UniformSampler(Sampler):
+
+    def __init__(
+            self,
+            rvs: RandomVariable | Sequence[RandomVariable],
+            condition: Condition = (),
+            rand: Random = random,
+    ):
+        condition: Sequence[Indicator] = check_condition(condition)
+
+        self._yield_f: YieldF
+        if isinstance(rvs, RandomVariable):
+            # a single rv
+            rvs = (rvs,)
+            self._yield_f = lambda x: x.item()
+        else:
+            # a sequence of rvs
+            self._yield_f = lambda x: x.tolist()
+
+        super().__init__(rvs, condition)
+
+        # Group condition indicators by `rv_idx`.
+        conditioned_rvs: MapSet[int, int] = MapSet()
+        for ind in condition:
+            conditioned_rvs.add(ind.rv_idx, ind.state_idx)
+
+        def get_possible_states(_rv: RandomVariable) -> List[int]:
+            condition_states: Optional[Set[int]] = conditioned_rvs.get(_rv.idx)
+            if condition_states is None:
+                return list(range(len(_rv)))
+            else:
+                return [state_idx for state_idx in range(len(_rv)) if state_idx not in condition_states]
+
+        possible_states: List[List[int]] = [
+            get_possible_states(rv)
+            for rv in self.rvs
+        ]
+
+        self._possible_states: List[List[int]] = possible_states
+        self._rand: Random = rand
+        self._state_dtype: DType = dtype_for_state_indexes(self.rvs)
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        possible_states = self._possible_states
+        yield_f = self._yield_f
+        rand = self._rand
+        state = np.zeros(len(possible_states), dtype=self._state_dtype)
+        while True:
+            for i, l in enumerate(possible_states):
+                state_idx = rand.randrange(0, len(l))
+                state[i] = l[state_idx]
+            yield yield_f(state)

ck/sampling/wmc_direct_sampler.py

@@ -0,0 +1,169 @@
+from typing import Collection, Iterator, Sequence
+
+import numpy as np
+
+from ck.pgm import Instance
+from ck.probability.probability_space import dtype_for_state_indexes
+from ck.program.program_buffer import ProgramBuffer
+from ck.program.raw_program import RawProgram
+from ck.sampling.sampler import Sampler
+from ck.sampling.sampler_support import SampleRV, YieldF, SamplerInfo
+from ck.utils.np_extras import NDArrayNumeric
+from ck.utils.random_extras import Random
+
+
+class WMCDirectSampler(Sampler):
+
+    def __init__(
+            self,
+            sampler_info: SamplerInfo,
+            raw_program: RawProgram,
+            rand: Random,
+    ):
+        super().__init__(sampler_info.rvs, sampler_info.condition)
+        self._yield_f: YieldF = sampler_info.yield_f
+        self._rand: Random = rand
+        self._program_buffer = ProgramBuffer(raw_program)
+        self._sample_rvs: Sequence[SampleRV] = tuple(sampler_info.sample_rvs)
+        self._chain_rvs: Sequence[SampleRV] = tuple(
+            sample_rv for sample_rv in sampler_info.sample_rvs if sample_rv.copy_index is not None)
+        self._state_dtype = dtype_for_state_indexes(self.rvs)
+        self._max_number_of_states: int = max((len(rv) for rv in self.rvs), default=0)
+        self._slots_1: Collection[int] = sampler_info.slots_1
+
+        # Set up the input slots to 0 or 1 to respect conditioning and initial Markov chain states.
+        slots: NDArrayNumeric = self._program_buffer.vars
+        for slot in sampler_info.slots_0:
+            slots[slot] = 0
+        for slot in sampler_info.slots_1:
+            slots[slot] = 1
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        yield_f = self._yield_f
+        rand = self._rand
+        sample_rvs = self._sample_rvs
+        chain_rvs = self._chain_rvs
+        slots_1 = self._slots_1
+        program_buffer = self._program_buffer
+        slots: NDArrayNumeric = program_buffer.vars
+
+        # Calling wmc() will give the weighted model count for the state of the current input slots.
+        def wmc() -> float:
+            return program_buffer.compute().item()
+
+        # Set up working memory buffers
+        states = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+        buff_slots = np.zeros(self._max_number_of_states, dtype=np.uintp)
+        buff_states = np.zeros(self._max_number_of_states, dtype=self._state_dtype)
+
+        while True:
+            # Consider all possible instantiations given the conditions, c, where the instantiations are ordered.
+            # Let awmc(i|c) be the accumulated WMC of the ith instantiation.
+            # We want to find the smallest instantiation i such that
+            #     rnd <= awmc(i|c)
+            # where rnd is in [0, 1) * wmc().
+
+            rnd: float = rand.random() * wmc()
+
+            for sample_rv in sample_rvs:
+                # Prepare to loop over random variable states.
+                # Keep track of the non-zero slots in buff_slots and buff_states.
+                num_possible_states: int = 0
+                for j, slot in enumerate(sample_rv.slots):
+                    if slots[slot] != 0:
+                        buff_slots[num_possible_states] = slot
+                        buff_states[num_possible_states] = j
+                        num_possible_states += 1
+
+                if num_possible_states == 0:
+                    raise RuntimeError('zero probability')
+
+                # Try each possible state of the current random variable.
+                # Once a state is selected, then the following is true:
+                #   states[rv_position] = state
+                #   m_prev_states[rv_position] = state
+                #   slots set up to include condition rv = state.
+                #   rnd is reduced to account for the states skipped.
+                #
+                # We can do this either by sequentially checking each state or by doing
+                # a binary search. Here we start with binary search then finish sequentially
+                # once the candidates size falls below 'THRESHOLD'.
+
+                # Binary search
+                THRESHOLD = 2
+                lo: int = 0
+                hi: int = num_possible_states
+                w_0_mark: int = 0
+                w: float = 0
+                while lo + THRESHOLD < hi:
+                    mid: int = (lo + hi) // 2
+
+                    for i in range(mid, hi):
+                        slots[buff_slots[i]] = 0
+
+                    w = wmc()
+                    w_0_mark = mid
+                    if w < rnd:
+                        # wmc() is too low, the desired state is >= buff_states[mid]
+                        for i in range(mid, hi):
+                            slots[buff_slots[i]] = 1
+                        lo = mid
+                    else:
+                        # wmc() is too high, the desired state is < buff_states[mid]
+                        hi = mid
+
+                # Now the state we want is between lo (inclusive) and hi (exclusive).
+                # Slots at least up to lo will be set to 1.
+
+                # clear top slots, lo and up.
+                for k in range(lo, num_possible_states):
+                    slots[buff_slots[k]] = 0
+
+                # Adjust rnd to account for lo > 0.
+                if lo == 0:
+                    # The chances of this case may be low, but if so, then
+                    # slots[m_buff_slots[lo]] = 0  which implies wmc() == 0,
+                    # so we can save a call to wmc().
+                    pass
+                elif w_0_mark == lo:
+                    # We can use the last wmc() call, stored in w.
+                    # This saves a call to wmc().
+                    rnd -= w
+                else:
+                    rnd -= wmc()
+
+                # Clear remaining slots
+                for k in range(0, lo):
+                    slots[buff_slots[k]] = 0
+
+                # Sequential search
+                k = lo
+                while k < hi:
+                    slot = buff_slots[k]
+                    slots[slot] = 1
+                    w = wmc()
+                    if rnd < w:
+                        break
+                    slots[slot] = 0
+                    rnd -= w
+                    k += 1
+
+                slot = buff_slots[k]
+                state = buff_states[k]
+                slots[slot] = 1
+                states[sample_rv.index] = state
+
+            yield yield_f(states)
+
+            # Reset the one slots for the next iteration.
+            for slot in slots_1:
+                slots[slot] = 1
+
+            # Copy chain pairs for next iteration.
+            # (This writes over any initial chain conditions from slots_1.)
+            for sample_rv in chain_rvs:
+                rv_slots = sample_rv.slots
+                prev_state_idx: int = states.item(sample_rv.copy_index)
+                for slot in rv_slots:
+                    slots[slot] = 0
+                slots[rv_slots[prev_state_idx]] = 1

ck/sampling/wmc_gibbs_sampler.py

@@ -0,0 +1,147 @@
+from typing import Collection, Iterator, Sequence, List
+
+import numpy as np
+
+from ck.pgm import Instance
+from ck.probability.probability_space import dtype_for_state_indexes
+from ck.program.program_buffer import ProgramBuffer
+from ck.program.raw_program import RawProgram
+from ck.sampling.sampler import Sampler
+from ck.sampling.sampler_support import SampleRV, YieldF, uniform_random_sample, SamplerInfo
+from ck.utils.np_extras import NDArrayStates, NDArrayFloat64
+from ck.utils.random_extras import Random, random_permute
+
+
+class WMCGibbsSampler(Sampler):
+
+    def __init__(
+            self,
+            sampler_info: SamplerInfo,
+            raw_program: RawProgram,
+            rand: Random,
+            skip: int,
+            burn_in: int,
+            pr_restart: float,
+    ):
+        super().__init__(sampler_info.rvs, sampler_info.condition)
+        self._yield_f: YieldF = sampler_info.yield_f
+        self._rand: Random = rand
+        self._program_buffer = ProgramBuffer(raw_program)
+        self._sample_rvs: List[SampleRV] = list(sampler_info.sample_rvs)
+        self._state_dtype = dtype_for_state_indexes(self.rvs)
+        self._slots_0: Collection[int] = sampler_info.slots_0
+        self._slots_1: Collection[int] = sampler_info.slots_1
+        self._skip: int = skip
+        self._burn_in: int = burn_in
+        self._pr_restart: float = pr_restart
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        sample_rvs: List[SampleRV] = self._sample_rvs
+        rand: Random = self._rand
+        yield_f: YieldF = self._yield_f
+        slots_0: Collection[int] = self._slots_0
+        slots_1: Collection[int] = self._slots_1
+        program_buffer: ProgramBuffer = self._program_buffer
+        skip: int = self._skip
+        burn_in: int = self._burn_in
+        pr_restart: float = self._pr_restart
+
+        # Allocate working memory
+        state = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+        prs: Sequence[NDArrayFloat64] = tuple(
+            np.zeros(len(sample_rv.slots), dtype=np.float64)
+            for sample_rv in sample_rvs
+        )
+
+        # Set an initial system state
+        uniform_random_sample(sample_rvs, slots_0, slots_1, program_buffer.vars, state, rand)
+
+        # Run a burn in
+        for i in range(burn_in):
+            self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+
+        if pr_restart <= 0:
+            # There is no possibility of a restart
+            if skip == 0:
+                while True:
+                    self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                    yield yield_f(state)
+            else:
+                while True:
+                    for _ in range(skip):
+                        self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                    self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                    yield yield_f(state)
+
+        else:
+            # There is the possibility of a restart
+            while True:
+                for _ in range(skip):
+                    self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                yield yield_f(state)
+                if rand.random() < pr_restart:
+                    # Set an initial system state
+                    uniform_random_sample(sample_rvs, slots_0, slots_1, program_buffer.vars, state, rand)
+
+                    # Run a burn in
+                    for i in range(burn_in):
+                        self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+
+    @staticmethod
+    def _next_sample_gibbs(
+            sample_rvs: List[SampleRV],
+            slots_1: Collection[int],
+            program_buffer: ProgramBuffer,
+            prs: Sequence[NDArrayFloat64],
+            state: NDArrayStates,
+            rand: Random
+    ) -> None:
+        """
+        Updates the states to a random system and reconfigures program inputs to match.
+        """
+        prog_in = program_buffer.vars
+        random_permute(sample_rvs, rand=rand)
+        for sample_rv in sample_rvs:
+            rv_slots = sample_rv.slots
+            index = sample_rv.index
+
+            rv_pr: NDArrayFloat64 = prs[index]
+            s: int = state.item(index)
+
+            candidates = []
+            for slot_state, slot in enumerate(rv_slots):
+                if slot in slots_1:
+                    candidates.append((slot_state, slot))
+            assert len(candidates) > 0
+
+            # Compute conditioned marginals for the current rv
+            prog_in[rv_slots[s]] = 0
+            for slot_state, slot in candidates:
+                prog_in[slot] = 1
+                rv_pr[slot_state] = program_buffer.compute()
+                prog_in[slot] = 0
+
+            # Pick a new state based on the conditional probabilities
+            total = np.sum(rv_pr)
+            if total == 0.0:
+                # No state of the current rv has a non-zero probability when
+                # conditioned on the other random variables states.
+                # Pick a random state form a uniform distribution.
+                i = rand.randrange(0, len(candidates))
+                candidate = candidates[i]
+                # update the states array and the wmc input
+                state[index] = candidate[0]
+                prog_in[candidate[1]] = 1
+            else:
+                # Pick a state, sampled from the marginal distribution
+                r = rand.random() * total
+                slot = None
+                slot_state = None
+                for slot_state, slot in candidates:
+                    if r <= rv_pr[slot_state]:
+                        break
+                    r -= rv_pr[slot_state]
+                # update the states array and the wmc input
+                state[index] = slot_state
+                prog_in[slot] = 1