compiled-knowledge 4.0.0a5__cp313-cp313-macosx_10_13_universal2.whl

ck/sampling/wmc_metropolis_sampler.py

@@ -0,0 +1,159 @@
+from typing import Collection, Iterator, List, 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)
+                    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)
+                    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)
+                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,113 @@
+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:
+                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,153 @@
+"""
+A module with extra iteration functions.
+"""
+from functools import reduce as _reduce
+from itertools import combinations, chain
+from operator import mul as _mul
+from typing import Iterable, Tuple, Iterator, Sequence, TypeVar
+
+_T = TypeVar('_T')
+
+
+def flatten(iterables: Iterable[Iterable[_T]]) -> Iterator[_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) -> Iterator:
+    """
+    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) -> Iterator[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) -> Iterator[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]) -> Iterator[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))
+    """
+    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))

ck/utils/map_list.py

@@ -0,0 +1,128 @@
+"""
+This module defines a class "MapList" for mapping keys to lists.
+"""
+from __future__ import annotations
+
+from typing import TypeVar, Generic, List, Dict, MutableMapping, KeysView, ValuesView, ItemsView, Iterable, Iterator
+
+_K = TypeVar('_K')
+_V = TypeVar('_V')
+
+
+class MapList(Generic[_K, _V], MutableMapping[_K, List[_V]]):
+    """
+    A MapList keeps a list for each key, unlike a dict which keeps only
+    a single element for each key.
+    """
+    __slots__ = ('_map', )
+
+    def __init__(self, *args, **kwargs):
+        self._map: Dict[_K, List[_V]] = {}
+        self.update(*args, extend=False, **kwargs)
+
+    def __str__(self) -> str:
+        return str(self._map)
+
+    def __repr__(self) -> str:
+        args = ', '.join(f'{key!r}:{val!r}' for key, val in self.items())
+        class_name = self.__class__.__name__
+        return f'{class_name}({args})'
+
+    def __len__(self) -> int:
+        return len(self._map)
+
+    def __bool__(self) -> bool:
+        return len(self) > 0
+
+    def __getitem__(self, key) -> List[_V]:
+        return self._map[key]
+
+    def __setitem__(self, key: _K, val: List[_V]):
+        if not isinstance(val, list):
+            class_name = self.__class__.__name__
+            raise RuntimeError(f'every {class_name} value must be a list')
+        self._map[key] = val
+
+    def __delitem__(self, key: _K):
+        del self._map[key]
+
+    def __iter__(self) -> Iterator[_K]:
+        return iter(self._map)
+
+    def __contains__(self, key: _K) -> bool:
+        return key in self._map
+
+    def update(self, *args, extend=False, **kwargs):
+        k: _K
+        v: List[_V]
+        if extend:
+            for k, v in dict(*args, **kwargs).items():
+                self.extend(k, v)
+        else:
+            for k, v in dict(*args, **kwargs).items():
+                self[k] = v
+
+    def keys(self) -> KeysView[_K]:
+        return self._map.keys()
+
+    def values(self) -> ValuesView[List[_V]]:
+        return self._map.values()
+
+    def items(self) -> ItemsView[_K, List[_V]]:
+        return self._map.items()
+
+    def get(self, key: _K, default=None):
+        """
+        Get the list corresponding to the given key.
+        If the key is not yet in the MapList then the
+        supplied default will be returned.
+        """
+        return self._map.get(key, default)
+
+    def get_list(self, key: _K) -> List[_V]:
+        """
+        Get the list corresponding to the given key.
+
+        This method will always return a list in the MapList, even if
+        it requires a new list being created.
+
+        Modifying the returned list affects this MapList object.
+        """
+        the_list = self._map.get(key)
+        if the_list is None:
+            the_list = []
+            self._map[key] = the_list
+        return the_list
+
+    def append(self, key: _K, item: _V):
+        """
+        Append the given item to the list identified by the given key.
+        """
+        self.get_list(key).append(item)
+
+    def extend(self, key: _K, items: Iterable[_V]):
+        """
+        Extend the given item to the list identified by the given key.
+        """
+        return self.get_list(key).extend(items)
+
+    def extend_map_list(self, map_list: MapList[_K, _V]):
+        """
+        Add all the keyed given items to the list identified by each key.
+        """
+        for key, items in map_list.items():
+            self.extend(key, items)
+
+    def clear(self):
+        """
+        Remove all items.
+        """
+        return self._map.clear()
+
+    def clear_empty(self):
+        """
+        Remove all empty values.
+        """
+        keys_to_remove = [key for key, value in self._map.items() if len(value) == 0]
+        for key in keys_to_remove:
+            del self._map[key]

ck/utils/map_set.py

@@ -0,0 +1,128 @@
+"""
+This module defines a class "MapSet" for mapping keys to sets.
+"""
+from __future__ import annotations
+
+from typing import TypeVar, Generic, Set, Dict, MutableMapping, Iterator, KeysView, ValuesView, ItemsView, Iterable
+
+_K = TypeVar('_K')
+_V = TypeVar('_V')
+
+
+class MapSet(Generic[_K, _V], MutableMapping[_K, Set[_V]]):
+    """
+    A MapSet keeps a set for each key, unlike a dict which keeps only
+    a single element for each key.
+    """
+    __slots__ = ('_map',)
+
+    def __init__(self, *args, **kwargs):
+        self._map: Dict[_K, Set[_V]] = {}
+        self.update(*args, add_all=False, **kwargs)
+
+    def __str__(self) -> str:
+        return str(self._map)
+
+    def __repr__(self) -> str:
+        args = ', '.join(f'{key!r}:{val!r}' for key, val in self.items())
+        class_name = self.__class__.__name__
+        return f'{class_name}({args})'
+
+    def __len__(self) -> int:
+        return len(self._map)
+
+    def __bool__(self) -> bool:
+        return len(self) > 0
+
+    def __getitem__(self, key: _K) -> Set[_V]:
+        return self._map[key]
+
+    def __setitem__(self, key: _K, val: Set[_V]):
+        if not isinstance(val, set):
+            class_name = self.__class__.__name__
+            raise RuntimeError(f'every {class_name} value must be a set')
+        self._map[key] = val
+
+    def __delitem__(self, key: _K):
+        del self._map[key]
+
+    def __iter__(self) -> Iterator[_K]:
+        return iter(self._map)
+
+    def __contains__(self, key: _K) -> bool:
+        return key in self._map
+
+    def update(self, *args, add_all=True, **kwargs):
+        k: _K
+        v: Set[_V]
+        if add_all:
+            for k, v in dict(*args, **kwargs).items():
+                self.add_all(k, v)
+        else:
+            for k, v in dict(*args, **kwargs).items():
+                self[k] = v
+
+    def keys(self) -> KeysView[_K]:
+        return self._map.keys()
+
+    def values(self) -> ValuesView[Set[_V]]:
+        return self._map.values()
+
+    def items(self) -> ItemsView[_K, Set[_V]]:
+        return self._map.items()
+
+    def get(self, key: _K, default=None):
+        """
+        Get the set corresponding to the given key.
+        If the key is not yet in the MapSet then the
+        supplied default will be returned.
+        """
+        return self._map.get(key, default)
+
+    def get_set(self, key: _K) -> Set[_V]:
+        """
+        Get the set corresponding to the given key.
+
+        This method will always return a set in the MapSet, even if
+        it requires a new set being created.
+
+        Modifying the returned set affects this MapSet object.
+        """
+        the_set = self._map.get(key)
+        if the_set is None:
+            the_set = set()
+            self._map[key] = the_set
+        return the_set
+
+    def add(self, key: _K, item: _V):
+        """
+        Add the given item to the set identified by the given key.
+        """
+        self.get_set(key).add(item)
+
+    def add_all(self, key: _K, items: Iterable[_V]):
+        """
+        Add all the given items to the set identified by the given key.
+        """
+        return self.get_set(key).update(items)
+
+    def add_map_set(self, map_set: MapSet[_K, _V]):
+        """
+        Add all the keyed given items to the set identified by each key.
+        """
+        for key, items in map_set.items():
+            self.add_all(key, items)
+
+    def clear(self):
+        """
+        Remove all items.
+        """
+        return self._map.clear()
+
+    def clear_empty(self):
+        """
+        Remove all empty values.
+        """
+        keys_to_remove = [key for key, value in self._map.items() if len(value) == 0]
+        for key in keys_to_remove:
+            del self._map[key]