compiled-knowledge 4.0.0a5__cp313-cp313-macosx_10_13_universal2.whl

ck/pgm_circuit/program_with_slotmap.py

@@ -0,0 +1,234 @@
+from typing import Tuple, Iterator, Sequence, Dict, Iterable
+
+from ck.pgm import RandomVariable, rv_instances, Instance, rv_instances_as_indicators, Indicator, ParamId
+from ck.pgm_circuit.slot_map import SlotMap, SlotKey
+from ck.probability.probability_space import Condition, check_condition
+from ck.program.program_buffer import ProgramBuffer
+from ck.utils.np_extras import NDArray, NDArrayNumeric
+
+
+class ProgramWithSlotmap:
+    """
+    A class for bundling a program buffer with a slot-map, where the slot-map maps keys
+    (e.g., random variable indicators) to program input slots.
+    """
+
+    def __init__(
+            self,
+            program_buffer: ProgramBuffer,
+            slot_map: SlotMap,
+            rvs: Sequence[RandomVariable],
+            precondition: Sequence[Indicator]
+    ):
+        """
+        Construct a ProgramWithSlotmap object.
+
+        Args:
+            program_buffer: is a ProgramBuffer object which is a compiled circuit with input and output slots.
+            slot_map: a maps from a slot_key to input slot of 'program'.
+            rvs: a sequence of rvs used for setting program input slots, each rv
+                has a length and rv[i] is a unique 'indicator' across all rvs.
+            precondition: conditions on rvs that are compiled into the program.
+
+        """
+        self._program_buffer: ProgramBuffer = program_buffer
+        self._slot_map: SlotMap = slot_map
+        self._rvs: Tuple[RandomVariable, ...] = tuple(rvs)
+        self._precondition: Sequence[Indicator] = precondition
+
+        if len(rvs) != len(set(rv.idx for rv in rvs)):
+            raise ValueError('duplicate random variables provided')
+
+        # Given rv = rvs[i], then _rvs_slots[i][state_idx] gives the slot for rv[state_idx].
+        self._rvs_slots: Tuple[Tuple[int, ...], ...] = tuple(tuple(self._slot_map[ind] for ind in rv) for rv in rvs)
+
+        # Given rv = rvs[i], then _indicator_map maps[rv[j]] = (i, slot), where slot is for indicator rv[j].
+        self._indicator_map: Dict[Indicator, Tuple[int, int]] = {
+            ind: (i, slot_map[ind])
+            for i, rv in enumerate(rvs)
+            for ind in rv
+        }
+
+    @property
+    def rvs(self) -> Sequence[RandomVariable]:
+        """
+        What are the random variables considered as 'inputs'.
+        """
+        return self._rvs
+
+    @property
+    def precondition(self) -> Sequence[Indicator]:
+        """
+        Condition on `self.rvs` that is compiled into the program.
+        """
+        return self._precondition
+
+    @property
+    def slot_map(self) -> SlotMap:
+        return self._slot_map
+
+    def instances(self, flip: bool = False) -> Iterator[Instance]:
+        """
+        Enumerate instances of the random variables.
+
+        Each instance is a tuples of state indexes, co-indexed with the given random variables.
+
+        The order is the natural index order (i.e., last random variable changing most quickly).
+
+        Args:
+            flip: if true, then first random variable changes most quickly.
+
+        Returns:
+            an iteration over tuples, each tuple holds state indexes
+            co-indexed with the given random variables.
+        """
+        return rv_instances(*self._rvs, flip=flip)
+
+    def instances_as_indicators(self, flip: bool = False) -> Iterator[Sequence[Indicator]]:
+        """
+        Enumerate instances of the random variables.
+
+        Each instance is a tuples of indicators, co-indexed with the given random variables.
+
+        The order is the natural index order (i.e., last random variable changing most quickly).
+
+        Args:
+            flip: if true, then first random variable changes most quickly.
+
+        Returns:
+            an iteration over tuples, each tuples holds random variable indicators
+            co-indexed with the given random variables.
+        """
+        return rv_instances_as_indicators(*self._rvs, flip=flip)
+
+    def compute(self) -> NDArrayNumeric:
+        """
+        Execute the program to compute and return the result. As per `ProgramBuffer.compute`.
+
+        Warning:
+            when returning an array, the array is backed by the program buffer memory, not a copy.
+        """
+        return self._program_buffer.compute()
+
+    def compute_conditioned(self, *condition: Condition) -> NDArrayNumeric:
+        """
+        Equivalent to:
+            self.set_condition(*condition)
+            return self.compute()
+        """
+        self.set_condition(*condition)
+        return self.compute()
+
+    @property
+    def results(self) -> NDArrayNumeric:
+        """
+        Get the results of the last computation.
+        As per `ProgramBuffer.results`.
+
+        Warning:
+            the array is backed by the program buffer memory, not a copy.
+        """
+        return self._program_buffer.results
+
+    @property
+    def vars(self) -> NDArrayNumeric:
+        """
+        Return the input variables as a numpy array.
+        As per `ProgramBuffer.vars`.
+
+        Warning:
+            writing to the returned array will write to the input slots of the program buffer.
+        """
+        return self._program_buffer.vars
+
+    def __setitem__(self, item: int | slice | SlotKey | Iterable[SlotKey], value: float) -> None:
+        """
+        Set one or more input slot values, identified by slot keys.
+        """
+        if isinstance(item, (int, slice)):
+            self._program_buffer[item] = value
+        elif isinstance(item, (Indicator, ParamId)):
+            self._program_buffer[self._slot_map[item]] = value
+        else:
+            # Assume its iterable
+            for i in item:
+                self[i] = value
+
+    def __getitem__(self, item: int | slice | SlotKey) -> NDArrayNumeric:
+        """
+        Get an input slot value, identified by a slot key.
+        """
+        if isinstance(item, (int, slice)):
+            return self._program_buffer[item]
+        elif isinstance(item, (Indicator, ParamId)):
+            return self._program_buffer[self._slot_map[item]]
+        else:
+            raise IndexError('unknown index type')
+
+    def set_condition(self, *condition: Condition) -> None:
+        """
+        Set the input slots of random variables to 1, except where implied to
+        0 according to the given conditions.
+
+        Specifically:
+            each slot corresponding to an indicator given condition will be set to 1;
+
+            if a random variable is mentioned in the given indicators, then all
+            slots for indicators for that random variable, except for slots corresponding
+            to an indicator given condition;
+
+            if a random variable is not mentioned in the given condition, that random variable
+            will have all its slots set to 1.
+        """
+        condition: Sequence[Indicator] = check_condition(condition)
+
+        ind_slot_groups = [[] for _ in self._rvs_slots]
+        for ind in condition:
+            rv_idx, slot = self._indicator_map[ind]
+            ind_slot_groups[rv_idx].append(slot)
+
+        slots: NDArray = self._program_buffer.vars
+        for rv_slots, ind_slots in zip(self._rvs_slots, ind_slot_groups):
+            if len(ind_slots) == 0:
+                # this rv _is not_ mentioned in the indicators - marginalise it
+                for slot in rv_slots:
+                    slots[slot] = 1
+            else:
+                # this rv _is_ mentioned in the indicators - we set the mentioned slots to 1 and others to 0.
+                for slot in rv_slots:
+                    slots[slot] = 0
+                for slot in ind_slots:
+                    slots[slot] = 1
+
+    def set_rv(self, rv: RandomVariable, *values: float | int) -> None:
+        """
+        Set the input values of a random variable.
+
+        Args:
+            rv: a random variable whose indicators are in the slot map.
+            values: list of values, assumes len(values) == len(rv).
+        """
+        for i in range(len(rv)):
+            self[rv[i]] = values[i]
+
+    def set_rvs_uniform(self, *rvs: RandomVariable) -> None:
+        """
+        Set the input values for each rv in rvs to 1 / len(rv).
+
+        Args:
+            rvs: a collection of random variable whose indicators are in the slot map.
+        """
+        for rv in rvs:
+            value = 1.0 / len(rv)
+            for ind in rv:
+                self[ind] = value
+
+    def set_all_rvs_uniform(self) -> None:
+        """
+        Set the input values for each rv in rvs to 1 / len(rv).
+        """
+        slots: NDArray = self._program_buffer.vars
+        for rv_slots in self._rvs_slots:
+            value = 1.0 / len(rv_slots)
+            for slot in rv_slots:
+                slots[slot] = value

ck/pgm_circuit/slot_map.py

@@ -0,0 +1,35 @@
+from typing import Protocol, Optional, overload, Iterable, Tuple
+
+from ck.pgm import Indicator, ParamId
+
+# Type of a slot map key.
+SlotKey = Indicator | ParamId
+
+
+class SlotMap(Protocol):
+    """
+    A slotmap is a protocol for mapping keys (indicators and
+    parameter ids) to slots in a ProgramBuffer.
+
+    A Python dict[SlotKey, int] implements the protocol.
+    """
+
+    def __len__(self) -> int:
+        ...
+
+    @overload
+    def get(self, slot_key: SlotKey, default: None) -> Optional[int]:
+        ...
+
+    @overload
+    def get(self, slot_key: SlotKey, default: int) -> int:
+        ...
+
+    def get(self, slot_key: SlotKey, default: Optional[int]) -> Optional[int]:
+        ...
+
+    def __getitem__(self, slot_key: SlotKey) -> int:
+        ...
+
+    def items(self) -> Iterable[Tuple[SlotKey, int]]:
+        ...

ck/pgm_circuit/support/compile_circuit.py

@@ -0,0 +1,83 @@
+from typing import Optional, Sequence
+
+from ck.circuit import CircuitNode, TmpConst, Circuit
+from ck.circuit_compiler import CircuitCompiler
+from ck.circuit_compiler.llvm_compiler import DataType, DEFAULT_TYPE_INFO, compile_circuit
+from ck.circuit_compiler import DEFAULT_CIRCUIT_COMPILER
+from ck.pgm_circuit import PGMCircuit
+from ck.program import RawProgram
+
+
+def compile_results(
+        pgm_circuit: PGMCircuit,
+        results: Sequence[CircuitNode],
+        const_parameters: bool,
+        compiler: CircuitCompiler = DEFAULT_CIRCUIT_COMPILER,
+) -> RawProgram:
+    """
+    Compile a circuit to a raw program that calculates the given result.
+
+    Raises:
+        ValueError: if not all nodes are from the same circuit.
+
+    Args:
+        pgm_circuit: The circuit (and PGM) that will be compiled to a program.
+        results: the result circuit nodes for the returned program.
+        const_parameters: if True then any circuit variable representing a parameter value will
+            be made 'const' in the resulting program.
+        compiler: function from circuit nodes to raw program.
+
+    Returns:
+        a compiled RawProgram.
+    """
+    circuit: Circuit = pgm_circuit.circuit_top.circuit
+    if const_parameters:
+        parameter_values = pgm_circuit.parameter_values
+        number_of_indicators = pgm_circuit.number_of_indicators
+        with TmpConst(circuit) as tmp:
+            for slot, value in enumerate(parameter_values, start=number_of_indicators):
+                tmp.set_const(slot, value)
+            raw_program: RawProgram = compiler(*results, circuit=circuit)
+    else:
+        raw_program: RawProgram = compiler(*results, circuit=circuit)
+
+    return raw_program
+
+
+def compile_param_derivatives(
+        pgm_circuit: PGMCircuit,
+        self_multiply: bool = False,
+        params_value: Optional[float | int] = 1,
+        data_type: DataType = DEFAULT_TYPE_INFO,
+) -> RawProgram:
+    """
+    Compile the circuit to a program for computing the partial derivatives of the parameters.
+    partial derivatives are co-indexed with pgm_circuit.parameter_values.
+
+    Typically, this will grow the circuit by the addition of circuit nodes to compute the derivatives.
+
+    This uses the LLVM circuit compiler.
+
+    Args:
+        pgm_circuit: The circuit (and PGM) that will be compiled to a program.
+        self_multiply: if true then each partial derivative df/dx will be multiplied by x.
+        params_value: if not None, then circuit vars representing parameters will be temporarily
+            set to this value for compiling the program. Default is 1.
+        data_type: What data type to use for arithmetic calculations. Either a DataType member or TypeInfo.
+    """
+    top: CircuitNode = pgm_circuit.circuit_top
+    circuit: Circuit = top.circuit
+
+    start_idx = pgm_circuit.number_of_indicators
+    end_idx = start_idx + pgm_circuit.number_of_parameters
+    param_vars = circuit.vars[start_idx:end_idx]
+    derivatives = circuit.partial_derivatives(top, param_vars, self_multiply=self_multiply)
+
+    if params_value is not None:
+        with TmpConst(circuit) as tmp:
+            tmp.set_const(param_vars, params_value)
+            raw_program: RawProgram = compile_circuit(*derivatives, circuit=circuit, data_type=data_type)
+    else:
+        raw_program: RawProgram = compile_circuit(*derivatives, circuit=circuit, data_type=data_type)
+
+    return raw_program

ck/pgm_circuit/target_marginals_program.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from typing import Optional, Tuple, List
+
+from ck.circuit import CircuitNode, Circuit, TmpConst
+from ck.pgm import RandomVariable
+from ck.pgm_circuit import PGMCircuit
+from ck.pgm_circuit.program_with_slotmap import ProgramWithSlotmap
+from ck.pgm_circuit.slot_map import SlotMap
+from ck.pgm_circuit.support.compile_circuit import compile_results
+from ck.probability.probability_space import check_condition, Condition
+from ck.program.program_buffer import ProgramBuffer
+from ck.program.raw_program import RawProgram
+from ck.utils.np_extras import NDArray
+
+
+class TargetMarginalsProgram(ProgramWithSlotmap):
+    """
+    """
+
+    def __init__(
+            self,
+            pgm_circuit: PGMCircuit,
+            target_rv: RandomVariable,
+            const_parameters: bool = True,
+    ):
+        """
+        Construct a TargetMarginalsProgram object.
+
+        Compile the given circuit for computing marginal probabilities over the states of 'target_var'.
+
+        Args:
+            pgm_circuit: The circuit representing a PGM.
+            target_rv: the random variable to compute marginals for.
+            const_parameters: if True then any circuit variable representing a parameter value will
+                be made 'const' in the resulting program.
+        """
+        top_node: CircuitNode = pgm_circuit.circuit_top
+        circuit: Circuit = top_node.circuit
+        slot_map: SlotMap = pgm_circuit.slot_map
+        input_rvs: List[RandomVariable] = list(pgm_circuit.rvs)
+
+        target_vars = [circuit.vars[slot_map[ind]] for ind in target_rv]
+        cct_outputs = circuit.partial_derivatives(top_node, target_vars)
+
+        # Remove the target rv from the input rvs.
+        target_index = input_rvs.index(target_rv)  # will throw if not found
+        del input_rvs[target_index]
+
+        with TmpConst(circuit) as tmp:
+            tmp.set_const(target_vars, 1)
+            raw_program: RawProgram = compile_results(
+                pgm_circuit=pgm_circuit,
+                results=cct_outputs,
+                const_parameters=const_parameters,
+            )
+
+        ProgramWithSlotmap.__init__(self, ProgramBuffer(raw_program), slot_map, input_rvs, pgm_circuit.conditions)
+
+        # additional fields
+        self._x_slots: List[List[int]] = [[slot_map[ind] for ind in rv] for rv in input_rvs]
+        self._y_size: int = raw_program.number_of_results
+        self._target_rv: RandomVariable = target_rv
+        self._number_of_indicators: int = pgm_circuit.number_of_indicators
+        self._z_cache: Optional[float] = None
+
+        # consistency check
+        assert (self._y_size == len(self._target_rv))
+
+        if not const_parameters:
+            # set the parameter slots
+            self.vars[pgm_circuit.number_of_indicators:] = pgm_circuit.parameter_values
+
+    @property
+    def target_rv(self) -> RandomVariable:
+        return self._target_rv
+
+    def map(self, condition: Condition = ()) -> Tuple[float, int]:
+        """
+        Return the maximum a posterior (MAP) state of the target variable.
+
+        Args:
+            condition: any conditioning indicators.
+
+        Returns:
+            (pr, state_idx) where
+            pr is the MAP probability
+            state_idx: is the MAP state index of `self.target_rv`.
+        """
+        self.set_condition(*check_condition(condition))
+        self.compute()
+        results: NDArray = self.results
+        z: float = results.sum()
+
+        max_p = -1
+        max_i = -1
+        for i in range(self._y_size):
+            p = results[i]
+            if p > max_p:
+                max_p = p
+                max_i = i
+
+        return max_p / z, max_i