compiled-knowledge 4.0.0a20__cp312-cp312-macosx_10_13_x86_64.whl

ck/program/program.py

@@ -0,0 +1,137 @@
+"""
+For more documentation on this module, refer to the Jupyter notebook docs/6_circuits_and_programs.ipynb.
+"""
+from typing import Callable, Sequence
+
+import numpy as np
+
+from ck.program.raw_program import RawProgram
+from ck.utils.np_extras import DTypeNumeric, NDArrayNumeric
+
+
+class Program:
+    """
+    A program represents an arithmetic a function from input values to output values.
+
+    Internally a `Program` wraps a `RawProgram` which is the object returned by a circuit compiler.
+
+    Every `Program` has a numpy `dtype` which defines the numeric data type for input and output values.
+    Typically, the `dtype` of a program is a C style double.
+    """
+
+    def __init__(self, raw_program: RawProgram):
+        self._raw_program = raw_program
+        self.__call = self.__get_call_method()
+
+    @property
+    def dtype(self) -> DTypeNumeric:
+        """
+        What is the C 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
+
+    def __call__(self, *args: float | int) -> NDArrayNumeric | int | float:
+        """
+        Call the compiled program.
+
+        Returns:
+            either a single value or a numpy array, depending
+            on the construction arguments.
+
+        Raises:
+            ValueError: if the given number of argument != self.number_of_vars.
+        """
+        # dynamically defined at construction time
+        return self.__call(*args)
+
+    def __call_single(self, *args: float) -> int | float:
+        """
+        Returns a single result of type self._dtype.value.dtype
+        """
+        return self._raw_program(args).item()
+
+    def __call_empty(self, *args: float | int) -> NDArrayNumeric:
+        """
+        Returns a numpy array result of dtype self.dtype
+        """
+        if len(args) != self.number_of_vars:
+            raise ValueError(f'incorrect number of arguments: expected {self.number_of_vars}, got {len(args)}')
+        np_out = np.zeros(0, dtype=self.dtype)
+        return np_out
+
+    def __call_multi(self, *args: float | int) -> NDArrayNumeric:
+        """
+        Returns a numpy array result of dtype self.dtype
+        """
+        return self._raw_program(args)
+
+    def __get_call_method(self) -> Callable:
+        """
+        Choose a call method based on self._number_of_results
+        """
+        match self.number_of_results:
+            case 0:
+                return self.__call_empty
+            case 1:
+                return self.__call_single
+            case _:
+                return self.__call_multi
+
+    def __getstate__(self):
+        """
+        Support for pickle.
+        """
+        return {
+            '_raw_program': self._raw_program,
+        }
+
+    def __setstate__(self, state):
+        """
+        Support for pickle.
+        """
+        self._raw_program = state['_raw_program']
+        self.__call = self.__get_call_method()

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,67 @@
+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  values,
+#     [1]: temporary working memory,
+#     [2]: output values.
+RawProgramFunction = Callable[[ct.POINTER, ct.POINTER, ct.POINTER], None]
+
+
+@dataclass
+class RawProgram:
+    """
+    A raw program is returned by a circuit compiler to provide execution of
+    the function defined by a compiled circuit.
+
+    A `RawProgram` is a `Callable` with the signature:
+
+
+    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 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