compiled-knowledge 4.0.0a20__cp312-cp312-macosx_11_0_arm64.whl

ck/sampling/wmc_gibbs_sampler.py

@@ -0,0 +1,153 @@
+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)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
+                    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)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
+                    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)
+                # We know the yield function will always provide either ints or Instances
+                # noinspection PyTypeChecker
+                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

ck/sampling/wmc_metropolis_sampler.py

@@ -0,0 +1,165 @@
+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, uniform_random_sample, SamplerInfo
+from ck.utils.np_extras import NDArrayStates, DTypeStates
+from ck.utils.random_extras import Random
+
+
+class WMCMetropolisSampler(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: Sequence[SampleRV] = tuple(sampler_info.sample_rvs)
+        self._state_dtype: DTypeStates = 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 = self._sample_rvs
+        rand = self._rand
+        yield_f = self._yield_f
+        slots_0 = self._slots_0
+        slots_1 = self._slots_1
+        program_buffer = self._program_buffer
+        slots = program_buffer.vars
+        skip = self._skip
+        burn_in = self._burn_in
+        pr_restart = self._pr_restart
+
+        # Allocate working memory
+        state: NDArrayStates = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+
+        # set up the input slots to respect conditioning
+        for slot in slots_0:
+            slots[slot] = 0
+        for slot in slots_1:
+            slots[slot] = 1
+
+        # Convert sample slots to possibles
+        # And map slots to states.
+        possibles = []
+        for sample_rv in sample_rvs:
+            rv_possibles = []
+            for slot_state, slot in enumerate(sample_rv.slots):
+                if slots[slot] == 1:
+                    rv_possibles.append((slot_state, slot))
+            possibles.append((sample_rv.index, sample_rv.slots, rv_possibles))
+
+        # Set an initial valid system state
+        w: float = self._init_sample_metropolis(state)
+
+        # Run a burn in
+        for i in range(burn_in):
+            w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+
+        if pr_restart <= 0:
+            # There is no possibility of a restart
+            if skip == 0:
+                while True:
+                    w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
+                    yield yield_f(state)
+            else:
+                while True:
+                    for _ in range(skip):
+                        w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                    w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
+                    yield yield_f(state)
+
+        else:
+            # There is the possibility of a restart
+            while True:
+                for _ in range(skip):
+                    w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                # We know the yield function will always provide either ints or Instances
+                # noinspection PyTypeChecker
+                yield yield_f(state)
+
+                if rand.random() < pr_restart:
+                    # Set an initial valid system state
+                    w = self._init_sample_metropolis(state)
+                    # Run a burn in
+                    for i in range(burn_in):
+                        w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+
+    def _init_sample_metropolis(self, state: NDArrayStates) -> float:
+        """
+        Initialises the states to a valid random system and configures program inputs to match.
+        """
+        sample_rvs = self._sample_rvs
+        rand = self._rand
+        slots_0 = self._slots_0
+        slots_1 = self._slots_1
+        program_buffer = self._program_buffer
+        slots = program_buffer.vars
+
+        while True:
+            uniform_random_sample(sample_rvs, slots_0, slots_1, slots, state, rand)
+            w: float = program_buffer.compute().item()
+            if w >= 0:
+                return w
+
+    @staticmethod
+    def _next_sample_metropolis(
+            possibles,
+            program_buffer: ProgramBuffer,
+            state,
+            cur_w: float,
+            rand: Random,
+    ) -> float:
+        """
+        Updates the states to a random system and reconfigures program inputs to match.
+        """
+        prog_in = program_buffer.vars
+
+        # Generate a proposal.
+        # randomly choose a random variable
+        i = rand.randrange(0, len(possibles))
+        idx, rv_slots, rv_possibles = possibles[i]
+        # keep track of the current state slot
+        cur_s = state[idx]
+        cur_s_slot = rv_slots[cur_s]
+        # randomly choose a possible state
+        i = rand.randrange(0, len(rv_possibles))
+        s, s_slot = rv_possibles[i]
+
+        # set up state and program to compute weight
+        prog_in[cur_s_slot] = 0
+        prog_in[s_slot] = 1
+
+        # calculate the weight and test it
+        new_w: float = program_buffer.compute().item()
+        if rand.random() * cur_w < new_w:
+            # accept
+            state[idx] = s
+            return new_w
+        else:
+            # reject: set state and program to what it was before
+            prog_in[s_slot] = 0
+            prog_in[cur_s_slot] = 1
+            return cur_w

ck/sampling/wmc_rejection_sampler.py

@@ -0,0 +1,115 @@
+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, uniform_random_sample, SamplerInfo
+from ck.utils.np_extras import NDArrayNumeric
+from ck.utils.random_extras import Random
+
+
+class WMCRejectionSampler(Sampler):
+
+    def __init__(
+            self,
+            sampler_info: SamplerInfo,
+            raw_program: RawProgram,
+            rand: Random,
+            z: 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: Sequence[SampleRV] = tuple(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
+
+        # Initialise fields for tracking max_w
+        self._w_max = None  # estimated maximum weight for any one world
+        self._w_not_seen = z  # z - w_seen
+        self._w_high = 0.0  # highest instance wight seen so far
+        self._samples = set()  # what samples have we seen
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        sample_rvs = self._sample_rvs
+        rand = self._rand
+        yield_f = self._yield_f
+        slots_0 = self._slots_0
+        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()
+
+        # Allocate working memory to store a possible world
+        state: NDArrayNumeric = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+
+        # Initialise w_max to w_max_marginal, if not done yet.
+        if self._w_max is None:
+            w_max_marginal = self._w_not_seen  # initially set to z, so a 'large' weight
+
+            # Set up the input slots to 0 or 1 to respect conditioning and initial Markov chain states.
+            for slot in slots_0:
+                slots[slot] = 0
+            for slot in slots_1:
+                slots[slot] = 1
+
+            # Loop over the rvs
+            for sample_rv in sample_rvs:
+                rv_slots = sample_rv.slots
+                max_for_rv = 0
+                # Set all rv slots to 0
+                for slot_state, slot in enumerate(rv_slots):
+                    slots[slot] = 0
+                back_to_one = []
+                # Loop over state of the rv.
+                for slot_state, slot in enumerate(rv_slots):
+                    if slot in slots_1:
+                        slots[slot] = 1
+                        w: float = wmc()
+                        max_for_rv = max(max_for_rv, w)
+                        slots[slot] = 0
+                        back_to_one.append(slot)
+                # Set rv slots back to 1 as needed (ready for next rv).
+                for slot in back_to_one:
+                    slots[slot] = 1
+
+                w_max_marginal = min(w_max_marginal, max_for_rv)
+
+            self._w_max = w_max_marginal
+
+        while True:
+            uniform_random_sample(sample_rvs, slots_0, slots_1, slots, state, rand)
+            w: float = wmc()
+
+            if rand.random() * self._w_max < w:
+                # We know the yield function will always provide either ints or Instances
+                # noinspection PyTypeChecker
+                yield yield_f(state)
+
+            # Update w_not_seen and w_high to adapt w_max.
+            # We don't bother tracking seen samples once w_not_seen and w_high
+            # are close enough, or we have tracked too many samples.
+            if self._samples is not None:
+                s = tuple(state)
+                if s not in self._samples:
+                    self._samples.add(s)
+                    self._w_not_seen -= w
+                    self._w_high = max(self._w_high, w)
+                    w_max_tracked = max(self._w_high, self._w_not_seen)
+                    self._w_max = min(w_max_tracked, self._w_max)
+
+                    # See if we should stop tracking samples.
+                    if (
+                            self._w_not_seen - self._w_high < 0.001  # w_not_seen and w_high are close enough
+                            or len(self._samples) > 1000000  # tracked too many samples
+                    ):
+                        self._samples = None

ck/utils/iter_extras.py

@@ -0,0 +1,163 @@
+"""
+A module with extra iteration functions.
+"""
+from functools import reduce as _reduce
+from itertools import combinations, chain, islice
+from operator import mul as _mul
+from typing import Iterable, Tuple, Sequence, TypeVar
+
+_T = TypeVar('_T')
+
+
+def flatten(iterables: Iterable[Iterable[_T]]) -> Iterable[_T]:
+    """
+    Iterate over the elements of an iterable of iterables.
+    """
+    return (elem for iterable in iterables for elem in iterable)
+
+
+def deep_flatten(iterables: Iterable) -> Iterable:
+    """
+    Iterate over the flattening of nested iterables.
+    """
+    for el in iterables:
+        if isinstance(el, Iterable) and not isinstance(el, str):
+            for sub in deep_flatten(el):
+                yield sub
+        else:
+            yield el
+
+
+def combos(list_of_lists: Sequence[Sequence[_T]], flip=False) -> Iterable[Tuple[_T, ...]]:
+    """
+    Iterate over all combinations of taking one element from each of the lists.
+
+    The order of results has the first element changing most rapidly.
+    For example, given [[1,2,3],[4,5],[6,7]], combos yields the following:
+        (1,4,6), (2,4,6), (3,4,6), (1,5,6), (2,5,6), (3,5,6),
+        (1,4,7), (2,4,7), (3,4,7), (1,5,7), (2,5,7), (3,5,7).
+
+    If flip, then the last changes most rapidly.
+    """
+    num = len(list_of_lists)
+    if num == 0:
+        yield ()
+        return
+    rng = range(num)
+    indexes = [0] * num
+    if flip:
+        start = num - 1
+        inc = -1
+        end = -1
+    else:
+        start = 0
+        inc = 1
+        end = num
+    while True:
+        yield tuple(list_of_lists[i][indexes[i]] for i in rng)
+        i = start
+        while True:
+            indexes[i] += 1
+            if indexes[i] < len(list_of_lists[i]):
+                break
+            indexes[i] = 0
+            i += inc
+            if i == end:
+                return
+
+
+def combos_ranges(list_of_lens: Sequence[int], flip=False) -> Iterable[Tuple[int, ...]]:
+    """
+    Equivalent to combos([range(l) for l in list_of_lens], flip).
+
+    The order of results has the first element changing most rapidly.
+    If flip, then the last changes most rapidly.
+    """
+    num = len(list_of_lens)
+    if num == 0:
+        yield ()
+        return
+    indexes = [0] * num
+    if flip:
+        start = num - 1
+        inc = -1
+        end = -1
+    else:
+        start = 0
+        inc = 1
+        end = num
+    while True:
+        yield tuple(indexes)
+        i = start
+        while True:
+            indexes[i] += 1
+            if indexes[i] < list_of_lens[i]:
+                break
+            indexes[i] = 0
+            i += inc
+            if i == end:
+                return
+
+
+def pairs(elements: Iterable[_T]) -> Iterable[Tuple[_T, _T]]:
+    """
+    Iterate over all possible pairs in the given list of elements.
+    """
+    return combinations(elements, 2)
+
+
+def sequential_pairs(elements: Sequence[_T]) -> Iterable[Tuple[_T, _T]]:
+    """
+    Iterate over sequential pairs in the given list of elements.
+    """
+    for i in range(len(elements) - 1):
+        yield elements[i], elements[i + 1]
+
+
+def powerset(iterable: Iterable[_T], min_size: int = 0, max_size: int = None) -> Iterable[Tuple[_T, ...]]:
+    """
+    powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)
+    """
+    if not isinstance(iterable, (list, tuple)):
+        iterable = list(iterable)
+    if min_size is None:
+        min_size = 0
+    if max_size is None:
+        max_size = len(iterable)
+    return chain.from_iterable(
+        combinations(iterable, size)
+        for size in range(min_size, max_size + 1)
+    )
+
+
+def unzip(xs: Iterable[Tuple[_T]]) -> Tuple[Iterable[_T]]:
+    """
+    Inverse function of zip.
+
+    E.g., a, b, c = unzip(zip(a, b, c))
+
+    Note that the Python type of `a`, `b`, and `c` may not be preserved, only
+    the contents, order and length are guaranteed.
+    """
+    return zip(*xs)
+
+
+def multiply(items: Iterable[_T], initial: _T = 1) -> _T:
+    """
+    Return the product of the given items.
+    """
+    return _reduce(_mul, items, initial)
+
+
+def first(items: Iterable[_T]) -> _T:
+    """
+    Return the first element of the iterable.
+    """
+    return next(iter(items))
+
+
+def take(iterable: Iterable[_T], n: int) -> Iterable[_T]:
+    """
+    Take the first n elements of the iterable.
+    """
+    return islice(iterable, n)