compiled-knowledge 4.0.0a5__cp313-cp313-macosx_10_13_universal2.whl

ck/program/program_buffer.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import ctypes as ct
+from typing import Sequence
+
+import numpy as np
+
+from ck.program.raw_program import RawProgram, RawProgramFunction
+from ck.utils.np_extras import DTypeNumeric, NDArrayNumeric
+
+
+class ProgramBuffer:
+    """
+    A ProgramBuffer wraps a RawProgram with pre-allocated input, tmp, and out buffers.
+    The buffers are numpy arrays.
+    """
+
+    def __init__(self, program: RawProgram):
+        self._raw_program: RawProgram = program
+
+        # Allocate the buffers
+        self._array_vars = np.zeros(self.number_of_vars, dtype=self.dtype)
+        self._array_tmps = np.zeros(self.number_of_tmps, dtype=self.dtype)
+        self._array_outs = np.zeros(self.number_of_results, dtype=self.dtype)
+
+        # Access the c-buffers
+        ptr_type = ct.POINTER(np.ctypeslib.as_ctypes_type(self.dtype))
+        self._c_array_vars = self._array_vars.ctypes.data_as(ptr_type)
+        self._c_array_tmps = self._array_tmps.ctypes.data_as(ptr_type)
+        self._c_array_outs = self._array_outs.ctypes.data_as(ptr_type)
+
+        # Keep a direct reference to the internal callable.
+        self._function: RawProgramFunction = program.function
+
+    def clone(self) -> ProgramBuffer:
+        """
+        Take a copy of this program buffer, with the same raw program
+        and same input and output values, but with new memory allocation
+        for the buffers.
+        """
+        clone = ProgramBuffer(self._raw_program)
+        clone[:] = self._array_vars
+        clone.results[:] = self._array_outs
+        return clone
+
+    @property
+    def raw_program(self) -> RawProgram:
+        """
+        What is the wrapped Program.
+        """
+        return self._raw_program
+
+    @property
+    def dtype(self) -> DTypeNumeric:
+        """
+        What is the numpy data type of values.
+        This is the same as numpy and ctypes `dtype`.
+        """
+        return self._raw_program.dtype
+
+    @property
+    def number_of_vars(self) -> int:
+        """
+        How many input values are there to the function.
+        Each input value relates to a circuit VarNode, as
+        per method `var_indices`.
+
+        Returns:
+            the number of input values.
+        """
+        return self._raw_program.number_of_vars
+
+    @property
+    def number_of_tmps(self) -> int:
+        """
+        How many temporary values are there to the function.
+
+        Returns:
+            the number of temporary values.
+        """
+        return self._raw_program.number_of_tmps
+
+    @property
+    def number_of_results(self) -> int:
+        """
+        How many output values are there from the function.
+
+        Returns:
+            the number of output values.
+        """
+        return self._raw_program.number_of_results
+
+    @property
+    def var_indices(self) -> Sequence[int]:
+        """
+        Get the circuit `VarNode.index` for each function input.
+
+        Returns:
+            a list of the circuit VarNode indices, co-indexed
+            with the function input values.
+        """
+        return self._raw_program.var_indices
+
+    @property
+    def vars(self) -> NDArrayNumeric:
+        """
+        Return the input variables as a numpy array.
+        Writing to the returned array will write to the input slots of the program buffer.
+
+        Warning:
+            the array is backed by the program buffer memory, not a copy.
+        """
+        return self._array_vars
+
+    @property
+    def results(self) -> NDArrayNumeric:
+        """
+        Return the results as a numpy array.
+
+        Warning:
+            the array is backed by the program buffer memory, not a copy.
+        """
+        return self._array_outs
+
+    def compute(self) -> NDArrayNumeric:
+        """
+        Compute and return the results, as per `self.results`.
+
+        Warning:
+            the array is backed by the program buffer memory, not a copy.
+        """
+        self._function(self._c_array_vars, self._c_array_tmps, self._c_array_outs)
+        return self._array_outs
+
+    def __setitem__(self, idx: int | slice, value: float) -> None:
+        """
+        Set the value of the indexed input variable.
+        """
+        self._array_vars[idx] = value
+
+    def __getitem__(self, idx: int | slice) -> NDArrayNumeric:
+        """
+        Get the value of the indexed input variable.
+        """
+        return self._array_vars[idx]
+
+    def __len__(self) -> int:
+        """
+        Number of input variables.
+        """
+        return len(self._array_vars)
+
+    def __getstate__(self):
+        """
+        Support for pickle.
+        """
+        return {
+            '_raw_program': self._raw_program,
+            '_array_vars': self._array_vars,
+            '_array_tmps': self._array_tmps,
+            '_array_outs': self._array_outs,
+        }
+
+    def __setstate__(self, state):
+        """
+        Support for pickle.
+        """
+        self._raw_program = state['_raw_program']
+        self._array_vars = state['_array_vars']
+        self._array_tmps = state['_array_tmps']
+        self._array_outs = state['_array_outs']
+
+        # Access the c-buffers
+        ptr_type = ct.POINTER(self.dtype)
+        self._c_array_vars = self._array_vars.ctypes.data_as(ptr_type)
+        self._c_array_tmps = self._array_tmps.ctypes.data_as(ptr_type)
+        self._c_array_outs = self._array_outs.ctypes.data_as(ptr_type)
+
+        # Keep a direct reference to the internal callable.
+        self._function: RawProgramFunction = self._raw_program.function

ck/program/raw_program.py

@@ -0,0 +1,61 @@
+from dataclasses import dataclass
+from typing import Callable, Sequence
+
+import numpy as np
+import ctypes as ct
+
+
+from ck.utils.np_extras import NDArrayNumeric, DTypeNumeric
+
+# RawProgramFunction is a function of three ctypes arrays, returning nothing.
+# Args:
+#     [0]: input parameter values,
+#     [1]: working memory,
+#     [2]: output values.
+RawProgramFunction = Callable[[ct.POINTER, ct.POINTER, ct.POINTER], None]
+
+
+@dataclass
+class RawProgram:
+    """
+    Fields:
+        function: is a function of three ctypes arrays, returning nothing.
+        dtype: the numpy data type of  the array values.
+        number_of_vars: the number of input parameter values (first function argument).
+        number_of_tmps: the number of working memory values (second function argument).
+        number_of_results: the number of result values (third function argument).
+        var_indices:  maps the index of inputs (from 0 to self.number_of_vars - 1) to the index
+            of the corresponding circuit var.
+    """
+
+    function: RawProgramFunction
+    dtype: DTypeNumeric
+    number_of_vars: int
+    number_of_tmps: int
+    number_of_results: int
+    var_indices: Sequence[int]
+
+    def __call__(self, in_vars: NDArrayNumeric | Sequence[int | float]) -> NDArrayNumeric:
+        """
+        Call the raw program as a function from an array to an array.
+        """
+        array_vars: NDArrayNumeric
+        if isinstance(vars, np.ndarray):
+            array_vars = in_vars
+        else:
+            array_vars = np.array(in_vars, dtype=self.dtype)
+        if array_vars.shape != (self.number_of_vars,):
+            raise ValueError(f'input array incorrect shape: got {array_vars.shape} expected ({self.number_of_vars},)')
+        if array_vars.dtype != self.dtype:
+            raise ValueError(f'input array incorrect dtype: got {array_vars.dtype} expected {self.dtype}')
+
+        array_tmps: NDArrayNumeric = np.zeros(self.number_of_tmps, dtype=self.dtype)
+        array_outs: NDArrayNumeric = np.zeros(self.number_of_results, dtype=self.dtype)
+
+        ptr_type = ct.POINTER(np.ctypeslib.as_ctypes_type(self.dtype))
+        c_array_vars = array_vars.ctypes.data_as(ptr_type)
+        c_array_tmps = array_tmps.ctypes.data_as(ptr_type)
+        c_array_outs = array_outs.ctypes.data_as(ptr_type)
+
+        self.function(c_array_vars, c_array_tmps, c_array_outs)
+        return array_outs

ck/sampling/forward_sampler.py

@@ -0,0 +1,211 @@
+"""
+This module defines a Sampler for Bayesian networks using
+the Forward Sampling method.
+"""
+import random
+from dataclasses import dataclass
+from typing import Optional, Iterator, Tuple, Sequence, Dict, List, Set
+
+import numpy as np
+
+from ck.pgm import PGM, RandomVariable, PotentialFunction, Instance, Indicator
+from ck.probability.probability_space import Condition, check_condition, dtype_for_state_indexes
+from ck.sampling.sampler import Sampler
+from ck.sampling.sampler_support import YieldF
+from ck.utils.np_extras import NDArrayStates, DTypeStates
+from ck.utils.random_extras import Random
+
+
+class ForwardSampler(Sampler):
+    """
+    A ForwardSampler operates directly on a Bayesian network PGM.
+    It does not compile the PGM but directly uses the CPTs.
+
+    It determines the parent random variables to sample (recursively)
+    and samples them in a 'forward' order.
+
+    Conditioning is implemented by rejecting samples incompatible with the condition,
+    thus if the probability of the condition is low, the sampler will reject
+    many samples, which may be slow.
+
+    When unconditioned, this sampler can be very efficient.
+    """
+
+    def __init__(
+            self,
+            pgm: PGM,
+            rvs: Optional[RandomVariable | Sequence[RandomVariable]] = None,
+            condition: Condition = (),
+            rand: Random = random,
+            check_is_bayesian_network: bool = True
+    ):
+        """
+        Construct a forward sampler for the given PGM.
+
+        Args:
+            pgm: is the model to be sampled. It is assumed to be a proper Bayesian network.
+            rvs: the list of random variables to sample; the yielded state vectors are
+                co-indexed with rvs; if None, then the WMC rvs are used; if rvs is a single
+                random variable, then single samples are yielded.
+            condition: is a collection of zero or more conditioning indicators.
+            rand: provides the stream of random numbers.
+            check_is_bayesian_network: is a Boolean flag. If true, then the
+                constructor will raise an exception if pgm.check_is_bayesian_network()
+                is false.
+        """
+        if check_is_bayesian_network and not pgm.check_is_bayesian_network():
+            raise ValueError('the given PGM is not a Bayesian network')
+
+        condition: Sequence[Indicator] = check_condition(condition)
+
+        self._yield_f: YieldF
+        if rvs is None:
+            # all input rvs
+            rvs = pgm.rvs
+            self._yield_f = lambda x: x.tolist()
+        elif rvs in pgm.rvs:
+            # 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)
+
+        # Create a map from rv_idx to its factor.
+        # This assumes a Bayesian network structure - one factor per rv.
+        rv_factors = {
+            factor.rvs[0].idx: factor
+            for factor in pgm.factors
+        }
+
+        # Infer sampling order of random variables.
+        # Infer the mapping from random variables to yielded samples.
+        # Get a _SampleRV for each rv in rvs, and for pre-requisites.
+        # This assumes a Bayesian network structure - no directed loops.
+        sample_rvs_map: Dict[int, _SampleRV] = {}  # map from rv index to _SampleRV object
+        sample_rvs_list: List[_SampleRV] = []  # sequence of _SampleRV objects for sampling order
+        output_sample_rvs: Sequence[_SampleRV] = tuple(
+            _get_sample_rv(rv.idx, rv_factors, sample_rvs_map, sample_rvs_list)
+            for rv in rvs
+        )
+
+        # Add constraints to sample rvs, based on conditions.
+        ind: Indicator
+        for ind in condition:
+            sample_rv: _SampleRV = _get_sample_rv(ind.rv_idx, rv_factors, sample_rvs_map, sample_rvs_list)
+            constraints = sample_rv.constraints
+            if constraints is None:
+                constraints = sample_rv.constraints = set()
+            constraints.add(ind.state_idx)
+
+        # Store needed data.
+        self._dtype: DTypeStates = dtype_for_state_indexes(pgm.rvs)
+        self._rand: Random = rand
+        self._sample_rvs_list: Sequence[_SampleRV] = sample_rvs_list
+        self._output_sample_rvs: Sequence[_SampleRV] = output_sample_rvs
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        yield_f = self._yield_f
+        output_sample_rvs = self._output_sample_rvs
+
+        # Allocate working memory for yielded result
+        instance: NDArrayStates = np.zeros(len(output_sample_rvs), dtype=self._dtype)
+
+        # Allocate working memory for sampled states
+        states: NDArrayStates = np.zeros(len(self._sample_rvs_list), dtype=self._dtype)
+
+        while True:
+            success = False
+            while not success:
+                success = self._set_sample(states)
+
+            for i, sample_rv in enumerate(output_sample_rvs):
+                instance[i] = states[sample_rv.val_idx]
+
+            yield yield_f(instance)
+
+    def _set_sample(self, states: NDArrayStates) -> bool:
+        """
+        Set the states array with a sample.
+        The states array is co-indexed with self._sample_rvs_list.
+        If the sample has zero probability (i.e. it should be rejected)
+        then False is returned, otherwise True is returned.
+        """
+        sample_rvs_list = self._sample_rvs_list
+        rand = self._rand
+
+        for sample_rv in sample_rvs_list:
+            state = sample_rv.sample(rand, states)
+
+            # Check that a sample was possible (i.e., positive conditioned
+            # probability for some state of `sample_rv`).
+            if state is None:
+                return False
+
+            # Check for possible rejection based on conditioning
+            if sample_rv.constraints is not None and state not in sample_rv.constraints:
+                return False
+
+            states[sample_rv.val_idx] = state
+        return True
+
+
+@dataclass
+class _SampleRV:
+    rv: RandomVariable
+    function: PotentialFunction
+    val_idx: int
+    parent_val_indexes: Sequence[int]
+    constraints: Optional[Set[int]]
+
+    def sample(self, rand: Random, states: NDArrayStates) -> Optional[int]:
+        """
+        Return a random state index for the random variable, using
+        the appropriate distribution over possible states,
+        conditioned on the parent states that can be found in 'states'.
+
+        The function returns None if there was no state with a non-zero probability.
+
+        Returns:
+            state index from this random variable, or None if no state with a non-zero probability.
+        """
+        num_states: int = len(self.rv)
+        function: PotentialFunction = self.function
+        parent_val_indexes: Sequence[int] = self.parent_val_indexes
+
+        x: float = rand.random()  # uniform variate in [0, 1)
+
+        # Get the cumulative CPD, conditioned on parent states.
+        parent_states: Tuple[int, ...] = tuple(states.item(i) for i in parent_val_indexes)
+        total: float = 0
+        for state in range(num_states):
+            total += function[(state,) + parent_states]
+            if x < total:
+                return state
+
+        if total <= 0:
+            # No state with positive probability
+            return None
+        else:
+            return num_states - 1
+
+
+def _get_sample_rv(rv_idx: int, rv_factors, sample_rvs_map, sample_rvs_list) -> _SampleRV:
+    """
+    Get a _SampleRV for the indexed random variable.
+    This may be called recursively for the parents of the nominated rv.
+    """
+    sample_rv = sample_rvs_map.get(rv_idx)
+    if sample_rv is None:
+        factor = rv_factors[rv_idx]
+        rv = factor.rvs[0]
+        parent_val_indexes = tuple(
+            _get_sample_rv(parent.idx, rv_factors, sample_rvs_map, sample_rvs_list).val_idx
+            for parent in factor[1:]
+        )
+        sample_rv = _SampleRV(rv, factor.function, len(sample_rvs_map), parent_val_indexes, None)
+        sample_rvs_map[rv_idx] = sample_rv
+        sample_rvs_list.append(sample_rv)
+    return sample_rv

ck/sampling/marginals_direct_sampler.py

@@ -0,0 +1,113 @@
+from typing import Collection, Iterator, Dict, 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 NDArray, NDArrayNumeric
+from ck.utils.random_extras import Random
+
+
+class MarginalsDirectSampler(Sampler):
+
+    def __init__(
+            self,
+            sampler_info: SamplerInfo,
+            raw_program: RawProgram,
+            rand: Random,
+            rv_idx_to_result_offset: Dict[int, int],
+    ):
+        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
+
+        self._marginals: Sequence[NDArrayNumeric] = tuple(
+            self._program_buffer.results[
+            rv_idx_to_result_offset[sample_rv.rv.idx]
+            :
+            rv_idx_to_result_offset[sample_rv.rv.idx] + len(sample_rv.rv)
+            ]
+            for sample_rv in sampler_info.sample_rvs
+        )
+        # Set up the input slots to 0 or 1 to respect conditioning and initial Markov chain states.
+        slots: NDArray = 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
+        program_buffer = self._program_buffer
+        slots: NDArray = program_buffer.vars
+        marginals = self._marginals
+        slots_1 = self._slots_1
+
+        # Set up working memory buffer
+        states = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+
+        def compute() -> float:
+            # Compute the program results based on the current input slot values.
+            # Return the WMC.
+            return program_buffer.compute().item(-1)
+
+        while True:
+            wmc: float = compute()
+            rnd: float = rand.random() * wmc
+
+            for sample_rv in sample_rvs:
+                index: int = sample_rv.index
+                if index > 0:
+                    # No need to execute the program on the first time through
+                    # as it was done just before entering the loop.
+                    wmc = compute()
+
+                rv_dist: NDArray = marginals[sample_rv.index]
+
+                rv_dist_sum: float = rv_dist.sum()
+                if rv_dist_sum <= 0:
+                    raise RuntimeError('zero probability')
+                rv_dist *= wmc / rv_dist_sum
+
+                state_index: int = -1
+                for i in range(len(sample_rv.rv)):
+                    w = rv_dist.item(i)
+                    if rnd < w:
+                        state_index = i
+                        break
+                    rnd -= w
+                assert state_index >= 0
+
+                for slot in sample_rv.slots:
+                    slots[slot] = 0
+                slots[sample_rv.slots[state_index]] = 1
+                states[index] = state_index
+
+            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/sampler.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from itertools import islice
+from typing import Sequence, Iterator
+
+from ck.pgm import RandomVariable, Instance, Indicator
+
+
+class Sampler(ABC):
+    """
+    A Sampler provides an unlimited series of samples for one or more random variables.
+    The random variables being sampled are provided as a tuple via the `rvs` property.
+
+    A Sampler will either iterate over Instance objects, where each instance is co-indexed
+    with `self.rvs`, or may iterate over single state indexes. Whether a Sampler iterates
+    over Instance objects or single state indexes is determined by the implementation.
+    If iterating over single state indexes, then `len(self.rvs) == 1`.
+    """
+    __slots__ = ('_rvs', '_condition')
+
+    def __init__(self, rvs: Sequence[RandomVariable], condition: Sequence[Indicator]):
+        """
+        Args:
+            rvs: a collection of the random variables being
+                sampled, co-indexed with each sample provided by `iter(self)`.
+            condition: condition on `rvs` that are compiled into the sampler.
+        """
+        self._rvs: Sequence[RandomVariable] = tuple(rvs)
+        self._condition: Sequence[Indicator] = tuple(condition)
+
+    @property
+    def rvs(self) -> Sequence[RandomVariable]:
+        """
+        What random variables are being sampled.
+
+        Returns:
+            the random variables being sampled, co-indexed with each sample from `iter(self)`.
+        """
+        return self._rvs
+
+    @property
+    def condition(self) -> Sequence[Indicator]:
+        """
+        Condition on `self.rvs` that are compiled into the sampler.
+        """
+        return self._condition
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        """
+        An unlimited series of samples from a random process.
+        Each sample is co-indexed with the random variables provided by `self.rvs`.
+        """
+        ...
+
+    def take(self, number_of_samples: int) -> Iterator[Instance] | Iterator[int]:
+        """
+        Take a limited number of samples from `iter(self)`.
+
+        Args:
+            number_of_samples: a limit on the number of samples to provide.
+        """
+        return islice(self, number_of_samples)