Skip to content

Commit

Permalink
Merge pull request #87 from simple-crypto/factor-graph
Browse files Browse the repository at this point in the history 
Factor graph: fix broken serialization and numerical bug in BP.
  • Loading branch information
@cassiersg
cassiersg authored Feb 20, 2023
2 parents 10a3e12 + d912f13 commit 496b7ce
Show file tree
Hide file tree
Showing 7 changed files with 296 additions and 4 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@ Changelog
Not released
------------

* Add ``vars`` and ``factors`` methods on `FactorGraph`.
* Add ``fg`` property on `BPState`.

v0.5.2 (2023/02/17)
-------------------

Expand Down
27 changes: 27 additions & 0 deletions src/scalib/attacks/factor_graph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,14 @@ def sanity_check(self, pub_assignment: ValsAssign, var_assignment: ValsAssign):
"""
self._inner.sanity_check(pub_assignment, var_assignment)

def vars(self) -> Sequence[str]:
"""Return the names of the variables in the graph."""
return self._inner.var_names()

def factors(self) -> Sequence[str]:
"""Return the names of the factors in the graph."""
return self._inner.factor_names()


class BPState:
"""Belief propagation state associated to a :class:`FactorGraph`.
Expand All @@ -164,8 +172,14 @@ def __init__(
):
if public_values is None:
public_values = dict()
self._fg = factor_graph
self._inner = factor_graph._inner.new_bp(nexec, public_values)

@property
def fg(self) -> FactorGraph:
"""The associated factor graph."""
self._fg

def set_evidence(self, var: str, distribution: Optional[npt.NDArray[np.float64]]):
r"""Sets prior distribution of a variable.
Expand All @@ -186,6 +200,19 @@ def set_evidence(self, var: str, distribution: Optional[npt.NDArray[np.float64]]
def bp_loopy(self, it: int, initialize_states: bool):
"""Runs belief propagation algorithm on the current state of the graph.
This is a shortcut for calls to :meth:`propagate_var` and :meth:`propagate_factor`. It is equivalent to:
.. code-block:: python
if initialize_states:
for var in self.fg.vars():
self.propagate_var(var)
for _ in range(it):
for factor in self.fg.factors():
self.propagate_factor(factor)
for var in self.fg.vars():
self.propagate_var(var)
Parameters
----------
it :
Expand Down
34 changes: 32 additions & 2 deletions src/scalib_ext/scalib-py/src/factor_graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,13 +58,15 @@ impl FactorGraph {
}

pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
Ok(PyBytes::new(py, &serialize(&self.inner.as_deref()).unwrap()).to_object(py))
let to_ser: Option<&sasca::FactorGraph> = self.inner.as_deref();
Ok(PyBytes::new(py, &serialize(&to_ser).unwrap()).to_object(py))
}

pub fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
match state.extract::<&PyBytes>(py) {
Ok(s) => {
self.inner = Some(Arc::new(deserialize(s.as_bytes()).unwrap()));
let deser: Option<sasca::FactorGraph> = deserialize(s.as_bytes()).unwrap();
self.inner = deser.map(Arc::new);
Ok(())
}
Err(e) => Err(e),
Expand Down Expand Up @@ -255,6 +257,34 @@ impl BPState {
self.get_inner_mut().propagate_factor_all(factor_id);
Ok(())
}
pub fn set_belief_from_var(
&mut self,
py: Python,
var: &str,
factor: &str,
distr: PyObject,
) -> PyResult<()> {
let edge_id = self.get_edge_named(var, factor)?;
let bp = self.get_inner_mut();
let distr = obj2distr(py, distr, bp.get_graph().edge_multi(edge_id))?;
bp.set_belief_from_var(edge_id, distr)
.map_err(|e| PyTypeError::new_err(e.to_string()))?;
Ok(())
}
pub fn set_belief_to_var(
&mut self,
py: Python,
var: &str,
factor: &str,
distr: PyObject,
) -> PyResult<()> {
let edge_id = self.get_edge_named(var, factor)?;
let bp = self.get_inner_mut();
let distr = obj2distr(py, distr, bp.get_graph().edge_multi(edge_id))?;
bp.set_belief_to_var(edge_id, distr)
.map_err(|e| PyTypeError::new_err(e.to_string()))?;
Ok(())
}
pub fn propagate_factor(
&mut self,
factor: &str,
Expand Down
14 changes: 14 additions & 0 deletions src/scalib_ext/scalib/src/sasca/belief_propagation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,20 @@ impl BPState {
pub fn get_belief_from_var(&self, edge: EdgeId) -> &Distribution {
&self.belief_from_var[edge]
}
pub fn set_belief_from_var(
&mut self,
edge: EdgeId,
belief: Distribution,
) -> Result<(), BPError> {
self.check_distribution(&belief, self.graph.edge_multi(edge))?;
self.belief_from_var[edge] = belief;
Ok(())
}
pub fn set_belief_to_var(&mut self, edge: EdgeId, belief: Distribution) -> Result<(), BPError> {
self.check_distribution(&belief, self.graph.edge_multi(edge))?;
self.belief_to_var[edge] = belief;
Ok(())
}
// Propagation type:
// belief to var -> var
// var -> belief to func
Expand Down
8 changes: 6 additions & 2 deletions src/scalib_ext/scalib/src/sasca/bp_compute.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -365,8 +365,12 @@ impl Distribution {
/// Normalize sum to one, and make values not too small
pub fn regularize(&mut self) {
self.for_each_ignore(|mut d, _| {
let norm_f = 1.0 / (d.sum() + MIN_PROBA * d.len() as f64);
d.mapv_inplace(|x| (x + MIN_PROBA) * norm_f);
let (sum, min) = d
.iter()
.fold((0.0f64, 0.0f64), |(sum, min), x| (sum + x, min.min(*x)));
let offset = -min + MIN_PROBA;
let norm_f = 1.0 / (sum + offset * d.len() as f64);
d.mapv_inplace(|x| (x + offset) * norm_f);
})
}
pub fn make_non_zero_signed(&mut self) {
Expand Down
6 changes: 6 additions & 0 deletions src/scalib_ext/scalib/src/sasca/factor_graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,12 @@ impl FactorGraph {
pub fn var_multi(&self, var: VarId) -> bool {
self.var(var).multi
}
pub fn factor_multi(&self, factor: FactorId) -> bool {
self.factor(factor).multi
}
pub fn edge_multi(&self, edge: EdgeId) -> bool {
self.factor_multi(self.edges[edge].factor)
}
pub fn range_vars(&self) -> impl Iterator<Item = VarId> {
(0..self.vars.len()).map(VarId::from_idx)
}
Expand Down
208 changes: 208 additions & 0 deletions tests/test_factorgraph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,37 @@ def make_distri(nc, n):
return normalize_distr(np.random.randint(1, 10000000, (n, nc)).astype(np.float64))


def test_copy_fg():
graph = """
NC 2
PROPERTY s1: x = a^b
VAR MULTI x
VAR MULTI a
VAR MULTI b
"""
graph = FactorGraph(graph)
graph2 = copy.deepcopy(graph)


def test_copy_bp():
graph = """
NC 2
PROPERTY s1: x = a^b
VAR MULTI x
VAR MULTI a
VAR MULTI b
"""
graph = FactorGraph(graph)
n = 5
bp_state = BPState(graph, n)
distri_a = make_distri(2, 5)
distri_b = make_distri(2, 5)
bp_state.set_evidence("a", distri_a)
bp_state.set_evidence("b", distri_b)
bp_state.bp_loopy(1, initialize_states=True)
bp_state2 = copy.deepcopy(bp_state)


def test_table():
"""
Test Table lookup
Expand Down Expand Up @@ -696,3 +727,180 @@ def test_cyclic():
g = FactorGraph(graph_multi_cyclic)
assert BPState(g, 1, {"p": np.array([0], dtype=np.uint32)}).is_cyclic() == False
assert BPState(g, 2, {"p": np.array([0, 0], dtype=np.uint32)}).is_cyclic() == True


def test_and_rounding_error_simple():
# simple reproduction of issue #86
factor_graph = """NC 16
VAR MULTI A
VAR MULTI B
VAR MULTI C
PROPERTY P: C = A & B
"""
priors = {
"A": [
1.0 + 2**-52,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
1.0,
0.0,
0.0,
0.0,
1.0,
],
"C": [
0.0,
2.666666666666667,
2.666666666666667,
0.0,
2.666666666666667,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
],
}
fg = FactorGraph(factor_graph)
bp = BPState(fg, 1)
for k, v in priors.items():
bp._inner.set_belief_from_var(k, "P", np.array([v]))
bp.propagate_factor("P")
assert (bp.get_belief_to_var("B", "P") >= 0.0).all()


def test_and_rounding_error_simple():
# test case of issue #86
factor_graph = """NC 16
VAR MULTI K0
VAR MULTI K1
VAR MULTI L1
VAR MULTI N1
VAR MULTI B
VAR MULTI C
VAR MULTI N0
VAR MULTI A
VAR MULTI L2
VAR MULTI L3
PUB SINGLE IV
PROPERTY P1: L1 = K0 ^ K1
PROPERTY P2: N1 = !K1
PROPERTY P3: B = N1 & IV
PROPERTY P4: C = B ^ K0
PROPERTY P5: N0 = !K0
PROPERTY P6: A = N0 & K1
PROPERTY P7: L2 = A ^ C
PROPERTY P8: L3 = K1 ^ C"""
priors = {
"A": [
0.0,
0.25,
0.25,
0.0,
0.25,
0.0,
0.0,
0.0,
0.25,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
],
"C": [
0.0,
0.0,
0.0,
0.1666666667,
0.0,
0.1666666667,
0.1666666667,
0.0,
0.0,
0.1666666667,
0.1666666667,
0.0,
0.1666666667,
0.0,
0.0,
0.0,
],
"L1": [
0.0,
0.0,
0.0,
0.1666666667,
0.0,
0.1666666667,
0.1666666667,
0.0,
0.0,
0.1666666667,
0.1666666667,
0.0,
0.1666666667,
0.0,
0.0,
0.0,
],
"L2": [
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
0.25,
0.0,
0.0,
0.0,
0.25,
0.0,
0.25,
0.25,
0.0,
],
"L3": [
0.0,
0.0,
0.0,
0.1666666667,
0.0,
0.1666666667,
0.1666666667,
0.0,
0.0,
0.1666666667,
0.1666666667,
0.0,
0.1666666667,
0.0,
0.0,
0.0,
],
}
fg = FactorGraph(factor_graph)
bp = BPState(fg, 1, public_values={"IV": 0xC})
for k, v in priors.items():
bp.set_evidence(k, distribution=np.array([v]))
bp.bp_loopy(5, initialize_states=False)
assert (bp.get_distribution("K0") >= 0.0).all()
assert (bp.get_distribution("K1") >= 0.0).all()

0 comments on commit 496b7ce

Please sign in to comment.