Repository: simple-crypto/SCALib (language: python)

Issue:
SASCA returns poor results for low noise cases
I am doing a SASCA on low-noise simulations, and noticed that for low noise cases (std < 0.25) SCAlib returns wrong results. Attached is code to reproduce and the factor graph 

![sasca](https://user-images.githubusercontent.com/6591138/205060420-89720345-2f4e-49fb-b376-1979b948cda8.png)

Above is the log of the rank vs the noise level (std). The blue line is the average rank of 10 experiments, and the box plot indicates the distribution of the ranks. Observe for noise=0.1, the distribution is spread. This was ran on the refactor-sasca branch.

```py 


import matplotlib.pyplot as plt
from scipy.stats import norm
from string import Template
import numpy as np
from scalib.attacks import FactorGraph, BPState

import re
class Circuit:
    def __init__(self, circ_file: str) -> None:
        self.instructions = Circuit.parse_circ(circ_file)
        self.inputs = [o for i in self.instructions for o in i[1:] if i[0] == "INPUT"]
        self.unmasked_inputs = [o for i in self.instructions for o in i[1:] if i[0] == 'INPUT_UNMASKED']
        self.outputs = [o for i in self.instructions for o in i[1:] if i[0] == "OUTPUT"]
        self.num_shares = [o for i in self.instructions for o in i[1:] if i[0] == "NUM_SHARES"]
        if self.num_shares:
            self.num_shares = int(self.num_shares[0])
        self.randomness = [
            o for i in self.instructions for o in i[1:] if i[0] == "RANDOM"
        ]
        self.variables = set([o for i in self.instructions for o in i[1:]])

    def evaluate(self, bitwidth=1):
        results = {k: 0 for k in self.variables}
        for i in self.inputs + self.randomness:
            results[i] = np.random.randint(2**bitwidth)
        for line in self.instructions:
            ins = line[0]
            operands = line[2:]
            res = line[1]
            match ins:
                case "XOR":
                    results[res] = results[operands[0]] ^ results[operands[1]]
                case "SHARES_XOR":
                    zxor = 0
                    [zxor := zxor ^ results[x] for x in operands]
                    results[res] = zxor
                case "AND":
                    results[res] = results[operands[0]] & results[operands[1]]
                case "NOT":
                    results[res] = results[operands[0]] ^ 0x1
                case "INPUT" | "OUTPUT" | "RANDOM" | "INPUT_UNMASKED" | "NUM_SHARES":
                    pass
                case _:
                    raise NotImplementedError(f"{ins} not implemented")
        return results

    @classmethod
    def parse_circ(cls, circ_file):
        return [
            re.findall(r"[\w']+", x)
            for x in [s.strip() for s in circ_file.split("\n")]
            if x
        ]

def circ2sascagraph(circuit, nc:int = 1) -> str:
    """Convert circ format to sasca graph format"""
    var_list = []
    properties = []
    graph = Template(f"""NC {2**nc}\n\n$vars\n\n$properties
    """)

    for line in circuit.instructions:
        ins = line[0]
        operands = line[1:]
        for o in operands:
            if o not in var_list:
                var_list.append(o)
        match ins:
            case 'SHARES_XOR':
                properties.append(f'PROPERTY {operands[0]} = {" ^ ".join(operands[1:])}')
            case 'XOR':
                properties.append(f'PROPERTY {operands[0]} = {operands[1]} ^ {operands[2]}')
            case 'AND':
                properties.append(f'PROPERTY {operands[0]} = {operands[1]} & {operands[2]}')
            case 'NOT':
                properties.append(f'PROPERTY {operands[0]} = !{operands[1]}')
            case _:
                pass
    
    return graph.substitute({
        "properties":'\n'.join(properties), 
        "vars": '\n'.join([f'VAR MULTI {v}' for v in var_list])})


def add_noise(evaluated_circuit: dict[str, int], noise: float):
            
    traces = {}
        
    for k,v in evaluated_circuit.items():
        if noise > 0:
            traces[k] = v + np.random.normal(0, noise) 
            p0 = norm.pdf(traces[k], 0, noise)
            p1 = norm.pdf(traces[k]-1, 0, noise)
            traces[k] = (p0/(p0+p1), p1/(p0+p1))
        else:
            traces[k] = (1, 0) if v == 0 else (0, 1)
    return traces
    

def HW(x):
    return np.sum(np.unpackbits(np.array([x], dtype=np.uint8)))

def add_hw_noise(evaluated_circuit: dict[str, int], noise: float, bitwidth=2):
    traces = {}
    for k,v in evaluated_circuit.items():
        lx = np.random.normal(HW(v), noise)
        p = np.zeros((2**bitwidth,))
        for i in range(2**bitwidth):
            p[i] = norm.pdf(lx, HW(i), noise)
        p = p/np.sum(p)
        traces[k] = p

    return traces

def compute_joint(posterior, inputs):
    joint = np.zeros((2**len(inputs),))

    for i in range(2**len(inputs)):
        ii = bin(i)[2:].zfill(len(inputs))
        j = 1
        c = len(inputs) - 1
        for n in inputs:
            j = j * posterior[n][-1][int(ii[c])]
            c -= 1
        joint[i] = round(j, 4)

    return joint
def compute_rank(input_val, joint):
    return np.log2((joint >= joint[input_val]).sum())
def eval_sasca(circuit, circ_state, priors, iterations=10):
    fg = FactorGraph(circ2sascagraph(circuit,nc=1))
    bp = BPState(fg, 1)
    for k,v in priors.items():
        if k in circuit.inputs:
            continue
        bp.set_evidence(k, distribution=np.array([
            [v[0], v[1]]
        ]))
    
    posteriors = {k: list() for k in circuit.unmasked_inputs}
    bp.bp_loopy(iterations)
    for k in circuit.unmasked_inputs:
        posteriors[k].append(bp.get_distribution(k)[0])
    input_val = int(''.join(
            [str(circ_state[o]) for o in sorted([o for o in circuit.unmasked_inputs], key=lambda a:a[1], reverse=True)]
    ),2)
    joint = compute_joint(posteriors, circuit.unmasked_inputs)
    rank = compute_rank(input_val, joint)

    return rank

if __name__ == "__main__":
    circ = """INPUT x0, x1, x2, x3, x4
OUTPUT c0, c1, c2, c3, c4
INPUT_UNMASKED x0, x1, x2, x3, x4

NOT(n0, x0)
NOT(n1, x1)
NOT(n2, x2)
NOT(n3, x3)
NOT(n4, x4)
AND(t0, x1, n0)
AND(t1, x2, n1)
AND(t2, x3, n2)
AND(t3, x4, n3)
AND(t4, x0, n4)
XOR(c0, x0, t1)
XOR(c1, x1, t2)
XOR(c2, x2, t3)
XOR(c3, x3, t4)
XOR(c4, x4, t0)"""


    circuit = Circuit(circ)
    noise_lvl = np.array([0.001, 0.1, 0.25, 0.5])
    
    bitwidth = 1
    num_experiments = 20
    sasca_results_rank = np.zeros((noise_lvl.shape[0], num_experiments))



    for n in range(num_experiments):
        circ_state = circuit.evaluate(bitwidth=bitwidth)
        priors = [add_hw_noise(circ_state, i, bitwidth=bitwidth) for i in noise_lvl]
        for i,p in enumerate(priors):
            sasca_results_rank[i][n] = eval_sasca(circuit, circ_state, p)
    
    plt.plot(noise_lvl, np.mean(sasca_results_rank, axis=1))
    plt.boxplot(sasca_results_rank.T, widths=0.3, positions=noise_lvl) 
    plt.ylabel('rank (log)')
    plt.xlabel('noise level (std)')
    plt.savefig('sasca.png')
```
```
NC 2

VAR MULTI x0
VAR MULTI x1
VAR MULTI x2
VAR MULTI x3
VAR MULTI x4
VAR MULTI c0
VAR MULTI c1
VAR MULTI c2
VAR MULTI c3
VAR MULTI c4
VAR MULTI n0
VAR MULTI n1
VAR MULTI n2
VAR MULTI n3
VAR MULTI n4
VAR MULTI t0
VAR MULTI t1
VAR MULTI t2
VAR MULTI t3
VAR MULTI t4

PROPERTY n0 = !x0
PROPERTY n1 = !x1
PROPERTY n2 = !x2
PROPERTY n3 = !x3
PROPERTY n4 = !x4
PROPERTY t0 = x1 & n0
PROPERTY t1 = x2 & n1
PROPERTY t2 = x3 & n2
PROPERTY t3 = x4 & n3
PROPERTY t4 = x0 & n4
PROPERTY c0 = x0 ^ t1
PROPERTY c1 = x1 ^ t2
PROPERTY c2 = x2 ^ t3
PROPERTY c3 = x3 ^ t4
PROPERTY c4 = x4 ^ t0
```


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 68d6c472..0969fe12 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,9 +8,9 @@
 dist/
 .tox/
 .coverage
-.cargo_build/
 docs/_build/
 htmlcov/
 src/scalib/version.py
 src/scalib/build_config.py
 .cargo_build*
+scalib*.whl
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8b95e8fb..ab992cc9 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -5,13 +5,24 @@ Changelog
 Not released
 ------------
 
-* Re-design `scalib.config` to handle more configuration in a single `Config` class. **Breaking change** to `scalib.config` and `scalib.config.threading`.
-* Smarter behavior for progress bars, unified configuration for progress bar. **Breaking change** to `scalib.attack.SASCAGraph`.
+* Deprecate ``SASCAGraph`` in favor of the new ``FactorGraph`` which provides:
+
+    - Overall much less restrictions on operations, including suport for:
+
+        + non-bijective tables
+        + multiple public operands
+        + NOT gate
+
+    - A ``sanity_check`` feature: verify that the graph is compatible with known value assignment (useful for debugging a graph).
+    - Optimized bitwise AND: ``O(nc*log(nc))`` instead of ``O(nc^2)``.
+
+* Re-design ``scalib.config`` to handle more configuration in a single ``Config`` class. **Breaking change** to ``scalib.config`` and ``scalib.config.threading``.
+* Smarter behavior for progress bars, unified configuration for progress bar. **Breaking change** to ``scalib.attack.SASCAGraph``.
 * Add accessors for the internal state of the LDA.
-* Introduce the `scalib.ScalibError` exception and remove `scalib.metrics.SnrError` (**Breaking change**).
+* Introduce the ``scalib.ScalibError`` exception and remove ``scalib.metrics.SnrError`` (**Breaking change**).
 * Improve error reporting in case of LDA solving error.
 * Allow all computations to be interrupted with Ctrl-C.
-* Fix deadlock when there is an error in large SNR computations (i.e., when n_vars*n_samples*n_traces > 2**33).
+* Fix deadlock when there is an error in large SNR computations (i.e., when ``n_vars*n_samples*n_traces > 2**33``).
 * Allow LDA to behave like simple pooled gaussian templates (#22).
 * Refresh build system (Tox version 4, improved CI), build by default with native machine optimizations.
 * Not crash anymore on non x86-64 CPUs (no CI for those yet).
diff --git a/docs/source/scalib.attacks.rst b/docs/source/scalib.attacks.rst
index fee7260f..588aad22 100644
--- a/docs/source/scalib.attacks.rst
+++ b/docs/source/scalib.attacks.rst
@@ -1,7 +1,11 @@
 Attacks
 =======
 
-SASCAGraph 
-----------
+.. autoclass:: scalib.attacks.FactorGraph
+   :members:
+
+.. autoclass:: scalib.attacks.BPState
+   :members:
+
 .. autoclass:: scalib.attacks.sascagraph.SASCAGraph
    :members:
diff --git a/setup.cfg b/setup.cfg
index a6c33a9f..863c14b1 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -20,7 +20,6 @@ python_requires = >=3.8
 install_requires =
         numpy~=1.19
         py-cpuinfo~=9.0
-        cpufeature~=0.2
 
 
 [options.packages.find]
diff --git a/src/scalib/attacks/__init__.py b/src/scalib/attacks/__init__.py
index b742b516..fcf45bfa 100644
--- a/src/scalib/attacks/__init__.py
+++ b/src/scalib/attacks/__init__.py
@@ -1,1 +1,2 @@
 from .sascagraph import SASCAGraph
+from .factor_graph import FactorGraph, BPState
diff --git a/src/scalib/attacks/factor_graph.py b/src/scalib/attacks/factor_graph.py
new file mode 100644
index 00000000..a49eec49
--- /dev/null
+++ b/src/scalib/attacks/factor_graph.py
@@ -0,0 +1,353 @@
+from typing import Sequence, Mapping, Union, Optional
+
+import numpy as np
+import numpy.typing as npt
+
+from scalib import _scalib_ext
+
+__all__ = ["FactorGraph", "BPState"]
+
+CstValue = Union[int, Sequence[int]]
+ValsAssign = Mapping[str, CstValue]
+
+
+class FactorGraph:
+    r"""FactorGraph allows to run Soft Analytical Side-Channel Attacks (SASCA).
+
+    A SASCA is based on a set of variables, on knowledge of
+    relationships between those variables and information about the values of
+    the variables.
+
+    Variables have values in a binary finite field of size `nc`, and are viewed
+    as bit vectors. Those values are represented as integers in [0, nc).
+    Variables can be qualified as `SINGLE` or `MULTI`, which relates
+    to the multiple-execution feature of `FactorGraph`: when performing a
+    SASCA, it is useful to acquire multiple execution traces. In these
+    executions, some variables stay the same (e.g. an encryption key), and other
+    change (e.g. plaintext, masked variables, etc.).
+    The `SINGLE` qualifier should be used for variables that remain the same
+    and `MULTI` for variables that change.
+    `FactorGraph` will then build a graph where each `MULTI` variable is
+    replicated `n` times (once for each execution), as well as all the
+    relationships that relate at least on `MULTI` variable.
+
+    Relationships between variables are bitwise XOR, bitwise AND, bitwise OR,
+    bitiwise negation, modular addition, modular multiplication and lookup
+    table. A lookup table can describe any function that maps a single variable
+    to another variable.
+    Description of `nc`, the variables, and the relationships is given in a
+    text format specified below.
+
+    Finally, variables are of two kinds: the public values and the variables.
+    The public values have a single known value, while variables have an
+    uncertain value modelled as a probability distribution.
+    The SASCA computes these distributions from the relationships encoded by
+    the graph, and from prior information.
+    The prior distributions are by default uniform, but this can be changed.
+
+    An attack attempting to recover the secret key byte `k` is shown below.
+
+    >>> # Describe and generate the SASCAGraph
+    >>> graph_desc = '''
+    ...     # Small unprotected Sbox example
+    ...     NC 256 # Graph over GF(256)
+    ...     TABLE sbox   # Sbox
+    ...     VAR SINGLE k # key (to recover !)
+    ...     PUB MULTI p  # plaintext (known)
+    ...     VAR MULTI x # Sbox input
+    ...     VAR MULTI y # Sbox output (whose leakage is targeted)
+    ...     PROPERTY x = k ^ p   # Key addition
+    ...     PROPERTY y = sbox[x] # Sbox lookup
+    ...     '''
+    >>> sbox = np.arange(256, dtype=np.uint32) # don't use this S-box ;)
+    >>> graph = FactorGraph(graph_desc, {"sbox": sbox})
+    >>> # n is the number of traces for our attack.
+    >>> n = 2
+    >>> plaintexts = np.array([0, 1], dtype=np.uint32)
+    >>> bp = BPState(graph, n, {"p": plaintexts})
+    >>> y_leakage = np.random.rand(2, 256) # this might come from an LDA
+    >>> y_leakage = y_leakage / y_leakage.sum(axis=1, keepdims=True)
+    >>> bp.set_evidence("y",y_leakage)
+    >>> # Solve graph
+    >>> bp.bp_loopy(it=3, initialize_states=True)
+    >>> # Get key distribution and derive key guess
+    >>> k_distri = bp.get_distribution("k")
+    >>> key_guess = np.argmax(k_distri)
+
+    By running a belief propagation algorithm (see [1]_), the distributions on all
+    the variables are updated based on their initial distributions. The
+    `SASCAGraph` can be solved by using `run_bp()`.
+
+    Notes
+    -----
+
+    **The graph description format** is a text line-oriented format (each line
+    of text is a statement) that describes a set of variables and relationships
+    between those variables.
+    The ordering of the statements is irrelevant and whitespace is irrelevant
+    except for newlines and around keywords. End-of-line comments start with
+    the `#` symbol.
+    The statements are:
+
+    - `NC <nc>`: specifies the field size (must be a power of two). There must
+      be one `NC` statement in the description.
+    - `PUB SINGLE|MULTI variable_name`: declares a public variable.  `variable_name`
+      is an identifier of the variable (allowed characters are letters, digits
+      and underscore). One of the qualifiers `SINGLE` or `MULTI` must be given.
+    - `VAR SINGLE|MULTI variable_name`: declares a variables.
+    - `PROPERTY w = x^y^z`: declares a bitwise XOR property. There can be any
+      number of operands.
+    - `PROPERTY x = x&y`: declares a bitwise AND property.
+    - `PROPERTY x = t[y]`: declares a LOOKUP property (`y` is the lookup of the
+      table `t` at index `y`). No public variable is allowed in this property.
+    - `PROPERTY x = !y`: declares a bitwise NOT property.
+      No public variable is allowed in this property.
+    - `TABLE` t = [0, 3, 2, 1]`: Declares a table that can be used in a LOOKUP.
+      The values provided in the table must belong to the interval [0, nc).
+      The initialization expression can be omitted from the graph description
+      (e.g. `TABLE t`) and be given with `tables` parameter.
+
+
+    **Note**: if the `MULTI` feature doesn't match your use-case, using only
+    `SINGLE` variables works (or set the number of execution to 1).
+
+
+    .. [1] "Soft Analytical Side-Channel Attacks". N. Veyrat-Charvillon, B.
+       Gérard, F.-X. Standaert, ASIACRYPT2014.
+
+    Parameters
+    ----------
+    graph: string
+        The graph description.
+    """
+
+    def __init__(
+        self,
+        graph_text: str,
+        tables: Optional[Mapping[str, npt.NDArray[np.uint32]]] = None,
+    ):
+        if tables is None:
+            tables = dict()
+        self._inner = _scalib_ext.FactorGraph(graph_text, tables)
+
+    def sanity_check(self, pub_assignment: ValsAssign, var_assignment: ValsAssign):
+        """Verify that the graph is compatible with example variable assignments.
+
+        If the graph is not compatible, raise a ``ValueError``.
+
+        Parameters
+        ----------
+        pub_assignment:
+            For each public variable its value for all test executions.
+        var_assignment:
+            For each non-public variable its value for all test executions.
+
+        Returns
+        -------
+        None
+        """
+        self._inner.sanity_check(pub_assignment, var_assignment)
+
+
+class BPState:
+    """Belief propagation state.
+
+    This is a stateful object on which belief propagation operations can be run.
+    See :class:`scalib.attacks.FactorGraph` for usage example.
+    """
+
+    def __init__(
+        self,
+        factor_graph: FactorGraph,
+        nexec: int,
+        public_values: Optional[ValsAssign] = None,
+    ):
+        if public_values is None:
+            public_values = dict()
+        self._inner = factor_graph._inner.new_bp(nexec, public_values)
+
+    def set_evidence(self, var: str, distribution: Optional[npt.NDArray[np.float64]]):
+        r"""Sets prior distribution of a variable.
+
+        Parameters
+        ----------
+        var :
+            Identifier of the variable to assign the distribution to.
+        distribution :
+            Distribution to assign. If `var` is SINGLE, must be of shape `(nc,)`.
+            If `var` is MULTI, must be of shape `(nexec,nc)`.
+            If None, sets the prior distribution to uniform.
+        """
+        if distribution is None:
+            self._inner.drop_evidence(var)
+        else:
+            self._inner.set_evidence(var, distribution)
+
+    def bp_loopy(self, it: int, initialize_states: bool):
+        """Runs belief propagation algorithm on the current state of the graph.
+
+        Parameters
+        ----------
+        it :
+            Number of iterations of belief propagation.
+        initialize_states:
+            Whether to update variable distributions before running the BP iterations.
+            Recommended after using :func:`BPState.set_evidence`.
+        """
+        if initialize_states:
+            self._inner.propagate_all_vars()
+        self._inner.propagate_loopy_step(it)
+
+    def bp_acyclic(
+        self,
+        dest: str,
+        *,
+        clear_intermediates: bool = True,
+        clear_evidence: bool = False,
+    ):
+        """Runs the non-loopy belief propagation algorithm on the current state of the graph.
+        This only works if the graph is not cyclic.
+
+        Parameters
+        ----------
+        dest:
+            Variable for which the belief propagation is computed.
+        clear_intermediates:
+            Drop the intermetidate distributions and beliefs that are computed.
+        clear_evidence:
+            Drop the evidence for the variables, once used in the algorithm.
+        """
+        self._inner.propagate_acyclic(dest, clear_intermediates, clear_evidence)
+
+    def get_distribution(self, var: str) -> Optional[npt.NDArray[np.float64]]:
+        r"""Returns the current distribution of a variable `var`.
+
+        Parameters
+        ----------
+        var : string
+            Identifier of the variable for which distribution must be returned.
+            Distribution cannot be obtained for public variables.
+
+        Returns
+        -------
+        distribution : array_like, f64
+            Distribution of `var`. If `var` is SINGLE, distribution has shape
+            `(nc)`. Else, it has shape `(n,nc)`.
+            If the variable has a uniform distribution, None may be returned
+            (but this is not guaranteed).
+        """
+        return self._inner.get_state(var)
+
+    def is_cyclic(self) -> bool:
+        """Test is the graph is cyclic."""
+        return self._inner.is_cyclic()
+
+    def set_distribution(
+        self, var: str, distribution: Optional[npt.NDArray[np.float64]]
+    ):
+        r"""Sets current distribution of a variable in the BP.
+
+        Parameters
+        ----------
+        var :
+            Identifier of the variable to assign the distribution to.
+        distribution :
+            Distribution to assign. If `var` is SINGLE, must be of shape `(nc,)`.
+            If `var` is MULTI, must be of shape `(nexec,nc)`.
+            If None, sets the distribution to uniform.
+        """
+        if distribution is None:
+            self._inner.drop_state(var)
+        else:
+            self._inner.set_state(var, distribution)
+
+    def get_belief_to_var(
+        self, var: str, factor: str
+    ) -> Optional[npt.NDArray[np.float64]]:
+        r"""Returns the current belief from factor to var.
+
+        Parameters
+        ----------
+        var : string
+            Identifier of the variable for which distribution must be returned.
+        factor : string
+            Identifier of the factor for which distribution must be returned.
+
+        Returns
+        -------
+        distribution : array_like, f64
+            Belief on the edge from `factor` to `var`. If `factor` is SINGLE, distribution has shape
+            `(nc)`. Else, it has shape `(n,nc)`.
+            If the belief is a uniform distribution, None may be returned
+            (but this is not guaranteed).
+        """
+        return self._inner.get_belief_to_var(var, factor)
+
+    def get_belief_from_var(
+        self, var: str, factor: str
+    ) -> Optional[npt.NDArray[np.float64]]:
+        r"""Returns the current belief from var to factor.
+
+        Parameters
+        ----------
+        var : string
+            Identifier of the variable for which distribution must be returned.
+        factor : string
+            Identifier of the factor for which distribution must be returned.
+
+        Returns
+        -------
+        distribution : array_like, f64
+            Belief on the edge from `var` to `factor`. If `factor` is SINGLE, distribution has shape
+            `(nc)`. Else, it has shape `(n,nc)`.
+            If the belief is a uniform distribution, None may be returned
+            (but this is not guaranteed).
+        """
+        return self._inner.get_belief_from_var(var, factor)
+
+    def propagate_var(self, var: str):
+        """Run belief propagation on variable var.
+
+        This fetches beliefs from adjacent factors, computes the var
+        distribution, and sends updated beliefs to all adjacent factors.
+
+        Parameters
+        ----------
+        var : string
+            Identifier of the variable.
+
+        """
+        return self._inner.propagate_var(var)
+
+    def propagate_factor(self, factor: str):
+        """Run belief propagation on the given factor.
+
+        This fetches beliefs from adjacent variables and sends updated beliefs
+        to all adjacent variables.
+
+        Parameters
+        ----------
+        var : string
+            Identifier of the variable.
+
+        """
+        return self._inner.propagate_factor_all(factor)
+
+    def debug(self):
+        """Debug-print the current state."""
+        s = []
+        s.append("VAR DISTRIBUTION")
+        for var in self._inner.graph().var_names():
+            s.append(f"\tVar {var}")
+            s.append(repr(self.get_distribution(var)))
+        s.append("FACTORS FROM VARS")
+        for factor in self._inner.graph().factor_names():
+            for var in self._inner.graph().factor_scope(factor):
+                s.append(f"\t{var} -> {factor}")
+                s.append(repr(self.get_belief_from_var(var, factor)))
+        s.append("FACTORS TO VARS")
+        for factor in self._inner.graph().factor_names():
+            for var in self._inner.graph().factor_scope(factor):
+                s.append(f"\t{factor} -> {var}")
+                s.append(repr(self.get_belief_to_var(var, factor)))
+        return "\n".join(s)
diff --git a/src/scalib/attacks/sascagraph.py b/src/scalib/attacks/sascagraph.py
index 3ea14413..28804a4b 100644
--- a/src/scalib/attacks/sascagraph.py
+++ b/src/scalib/attacks/sascagraph.py
@@ -1,4 +1,6 @@
+import copy
 from functools import reduce
+import math
 
 import numpy as np
 
@@ -6,10 +8,30 @@
 from scalib.config import get_config
 import scalib.utils
 
+__all__ = ["SASCAGraph"]
+
+_NODE_FN = {
+    "AND": lambda a, b: a & b,
+    "XOR": lambda a, b: a ^ b,
+    "ADD": lambda a, b: (a + b) % self.nc_,
+    "MUL": lambda a, b: (a * b) % self.nc_,
+}
+
+
+def _node_reduce(node, values):
+    l = [v for v in values if v is not None]
+    if l:
+        return reduce(_NODE_FN[node], l)
+    else:
+        return None
+
 
 class SASCAGraph:
     r"""SASCAGraph allows to run Soft Analytical Side-Channel Attacks (SASCA).
 
+     .. deprecated:: 0.5.0
+         Use :class:`scalib.attacks.FactorGraph` instead.
+
     A `SASCAGraph` attack is based on a set of variables, on knowledge of
     relationships between those variables and information about the values of
     the variables.
@@ -95,6 +117,8 @@ class SASCAGraph:
       exactly two operands, with at most one public operand.
     - `LOOKUP x = t[y]`: declares a LOOKUP property (`y` is the lookup of the
       table `t` at index `y`). No public variable is allowed in this property.
+    - `LOOKUP x = !y`: declares a bitwise NOT property.
+      No public variable is allowed in this property.
     - `TABLE` t = [0, 3, 2, 1]`: Declares a table that can be used in a LOOKUP.
       The values provided in the table must belong to the interval [0, nc).
       The initialization expression can be omitted from the graph description
@@ -118,6 +142,16 @@ class SASCAGraph:
     """
 
     def __init__(self, graph, n):
+        import warnings
+
+        warnings.simplefilter("once", DeprecationWarning)  # turn off filter
+        warnings.warn(
+            "SASCAGraph is deprecated and will be removed in a next release. Use FactorGraph instead.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        warnings.simplefilter("default", DeprecationWarning)  # reset filter
+
         self.n_ = n
         self.solved_ = False
 
@@ -130,6 +164,48 @@ def __init__(self, graph, n):
         self.var_ = self.graph.var
         self.publics_ = {}
 
+    def sanity_check(self, var_assignment):
+        """Verify that the graph is compatible with example variable assignments.
+
+        If the graph is not compatible, raise a SACAGraphError.
+
+        Parameters
+        ----------
+        var_assignment: dict[var_name -> np.array (int)]
+            For each non-public variable its value for all test executions.
+        """
+        for k in self.var_:
+            if k not in var_assignment:
+                raise ValueError(f"Missing value for variable {k}.")
+        for k in self.publics_:
+            if k in var_assignment:
+                raise ValueError(f"Value given for public variable {k}.")
+
+        def get_var_values(var):
+            if var in self.publics_:
+                return self.publics_[var]
+            else:
+                if self.var_[var]["para"]:
+                    return var_assignment[var]
+                else:
+                    return var_assignment[var] * np.ones((self.n_,), dtype=np.uint32)
+
+        for property in self.properties_:
+            inp_values = [get_var_values(inp) for inp in property["inputs"]]
+            if property["property"] == "LOOKUP":
+                output = self.tables_[property["tab"]][inp_values[0]]
+            elif property["property"] == "NOT":
+                output = (self.nc_ - 1) ^ inp_values[0]
+            elif property["property"] in _NODE_FN:
+                output = _node_reduce(property["property"], inp_values)
+            else:
+                assert False
+            output_val = get_var_values(property["output"])
+            if (output != output_val).any():
+                raise SASCAGraphError(
+                    f"Property {property} fail, expected: {output}, given: {output_val}."
+                )
+
     def set_init_distribution(self, var, distribution):
         r"""Sets initial distribution of a variables.
 
@@ -220,10 +296,6 @@ def set_table(self, table, values):
             raise ValueError("Table has wrong shape")
         elif values.dtype != np.uint32:
             raise ValueError("Table must be np.uint32")
-        elif set(values) != set(range(self.nc_)):
-            raise ValueError(
-                "In current implementation, table is not a bijection over the set of values [0, nc)."
-            )
         self.tables_[table] = values
 
     def run_bp(self, it):
@@ -261,18 +333,11 @@ def _check_fully_init(self):
 
     def _init_graph(self):
         self._check_fully_init()
-        # mapping to Rust functions
-        AND = 0
-        XOR = 1
-        XOR_CST = 2
-        LOOKUP = 3
-        AND_CST = 4
-        ADD = 5
-        ADD_CST = 6
-        MUL = 7
-        MUL_CST = 8
-        property_map_binary = {"AND": (AND, AND_CST), "MUL": (MUL, MUL_CST)}
-        property_map_nary = {"XOR": (XOR, XOR_CST), "ADD": (ADD, ADD_CST)}
+        unary_properties = [
+            "NOT",
+        ]
+        binary_properties = ["AND", "MUL"]
+        nary_properties = ["XOR", "ADD"]
 
         # edge id
         self.edge_ = 0
@@ -285,14 +350,10 @@ def _init_graph(self):
         for property in self.properties_:
             if property["output"] in self.publics_:
                 raise ValueError(
-                    "In current implementation public vars can only be ^ or & operands.\n"
+                    "In current implementation public vars can only be operands.\n"
                     + "Cannot assign "
                     + property["output"]
                 )
-            if len([inp for inp in property["inputs"] if inp in self.publics_]) > 1:
-                raise ValueError(
-                    "In current implementation there can only be one public operand."
-                )
             for inp in property["inputs"]:
                 if inp in self.publics_ and property["property"] == "LOOKUP":
                     raise ValueError(
@@ -319,85 +380,41 @@ def _init_graph(self):
                 )
             )
 
+            property["func"] = property["property"]
+
             if property["property"] == "LOOKUP":
-                property["func"] = LOOKUP
                 # get the table into the function
                 property["table"] = self.tables_[property["tab"]]
                 # set edge to input and output
                 self._share_edge(property, property["output"])
                 self._share_edge(property, property["inputs"][0])
-
-            elif property["property"] in property_map_binary.keys():
-                op, op_cst = property_map_binary[property["property"]]
-                # If no public inputs
-                if all(x in self.var_ for x in property["inputs"]):
-                    property["func"] = op
-                    self._share_edge(property, property["output"])
-                    for i in property["inputs"]:
-                        self._share_edge(property, i)
-
-                # if and with public input
-                elif len(property["inputs"]) == 2:
-                    # OP with one public
-                    property["func"] = op_cst
-
-                    self._share_edge(property, property["output"])
-
-                    # which of both inputs is public
-                    if property["inputs"][0] in self.var_:
-                        i = (0, 1)
-                    elif property["inputs"][1] in self.var_:
-                        i = (1, 0)
-                    else:
-                        assert False
-
-                    # share edge with non public input
-                    self._share_edge(property, property["inputs"][i[0]])
-
-                    # merge public input in the property
-                    property["values"] = self.publics_[property["inputs"][i[1]]]
-
-            elif property["property"] in property_map_nary.keys():
-                op, op_cst = property_map_nary[property["property"]]
-                # if no inputs are public
-                if all(x in self.var_ for x in property["inputs"]):
-                    # OP with no public
-                    property["func"] = op
-
-                    # share edge with the output
-                    self._share_edge(property, property["output"])
-
-                    # share edge with all the inputs
-                    for i in property["inputs"]:
-                        self._share_edge(property, i)
-
-                elif len(property["inputs"]) == 2:
-                    # OP with one public
-                    property["func"] = op_cst
-
-                    # which of both inputs is public
-                    if property["inputs"][0] in self.var_:
-                        i = (0, 1)
-                    elif property["inputs"][1] in self.var_:
-                        i = (1, 0)
-                    else:
-                        assert False
-
-                    # share edges with variables
-                    self._share_edge(property, property["output"])
-                    self._share_edge(property, property["inputs"][i[0]])
-
-                    # merge public into the property
-                    property["values"] = self.publics_[property["inputs"][i[1]]]
-
-                else:
-                    key = property["property"]
-                    raise ValueError(
-                        f"{key} must have two operands when one operand is public."
+            elif property["property"] in unary_properties:
+                assert len(property["inputs"]) == 1
+                self._share_edge(property, property["output"])
+                self._share_edge(property, property["inputs"][0])
+            elif (
+                property["property"] in binary_properties
+                or property["property"] in nary_properties
+            ):
+                if (
+                    property["property"] in binary_properties
+                    and len(property["inputs"]) != 2
+                ):
+                    raise SASCAGraphError(
+                        f"Too many inputs. output: {property['output']}"
                     )
-
+                property["values"] = _node_reduce(
+                    property["property"],
+                    (self.publics_.get(x) for x in property["inputs"]),
+                )
+                self._share_edge(property, property["output"])
+                for inp in property["inputs"]:
+                    if inp in self.var_:
+                        self._share_edge(property, inp)
+                if any(inp in self.publics_ for inp in property["inputs"]):
+                    property["func"] = property["property"] + "CST"
             else:
-                assert False, "Property must be either LOOKUP, AND or XOR."
+                assert False, "Property non-implemented."
 
         for v in self.var_:
             v = self.var_[v]
@@ -429,6 +446,16 @@ def __init__(self, description):
         self._build_var_set()
         self._build_tables()
         self._build_properties()
+        self._check_nc_bitwise()
+
+    def _check_nc_bitwise(self):
+        if bin(self.nc).count("1") == 1:
+            return
+        if any(prop["property"] in ("NOT", "XOR", "AND") for prop in self.properties):
+            self.errors.append(
+                "Use of bitwise operators with NC not a power of 2 is not supported."
+            )
+            self._raise_errors()
 
     def _build_properties(self):
         for prop_kind, res, inputs in self.prop_decls:
@@ -476,8 +503,8 @@ def _get_nc(self):
             self.errors.append("NC appears multiple times, can only appear once.")
         elif len(self.nc_decls) == 0:
             self.errors.append("NC not declared.")
-        elif self.nc_decls[0] not in range(1, 2**16 + 1):
-            self.errors.append("NC not in admissible range [1, 2^16].")
+        elif self.nc_decls[0] not in range(1, 2**32 + 1):
+            self.errors.append("NC not in admissible range [1, 2^32].")
         else:
             self.nc = self.nc_decls[0]
         self._raise_errors()
@@ -535,6 +562,12 @@ def _parse_sasca_property(self, rem):
             inputs = prop.split("&")
             if len(inputs) != 2:
                 raise SASCAGraphError("Wrong number of & operands: must be 2.")
+        elif "!" in prop:
+            prop_kind = "NOT"
+            inputs = prop.split("!")
+            if len(inputs) != 2 or inputs[0] != "":
+                raise SASCAGraphError("Wrong number of ! operands: must be 1 (post).")
+            inputs = [inputs[1]]
         elif "[" in prop and "]" in prop:
             prop_kind = "LOOKUP"
             tab, in_ = prop.split("[")
diff --git a/src/scalib_ext/Cargo.lock b/src/scalib_ext/Cargo.lock
index 1d4244f0..ede0b48f 100644
--- a/src/scalib_ext/Cargo.lock
+++ b/src/scalib_ext/Cargo.lock
@@ -2,6 +2,15 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "ahash"
+version = "0.3.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8fd72866655d1904d6b0997d0b07ba561047d070fbe29de039031c641b61217"
+dependencies = [
+ "const-random",
+]
+
 [[package]]
 name = "aho-corasick"
 version = "0.7.20"
@@ -38,6 +47,15 @@ dependencies = [
  "derive_arbitrary",
 ]
 
+[[package]]
+name = "ariadne"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1cb2a2046bea8ce5e875551f5772024882de0b540c7f93dfc5d6cf1ca8b030c"
+dependencies = [
+ "yansi",
+]
+
 [[package]]
 name = "atomic"
 version = "0.5.1"
@@ -64,6 +82,15 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
 
+[[package]]
+name = "bincode"
+version = "1.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "bindgen"
 version = "0.58.1"
@@ -183,6 +210,15 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
+[[package]]
+name = "chumsky"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8d02796e4586c6c41aeb68eae9bfb4558a522c35f1430c14b40136c3706e09e4"
+dependencies = [
+ "ahash",
+]
+
 [[package]]
 name = "clang-sys"
 version = "1.4.0"
@@ -232,6 +268,28 @@ dependencies = [
  "windows-sys 0.42.0",
 ]
 
+[[package]]
+name = "const-random"
+version = "0.1.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "368a7a772ead6ce7e1de82bfb04c485f3db8ec744f72925af5735e29a22cc18e"
+dependencies = [
+ "const-random-macro",
+ "proc-macro-hack",
+]
+
+[[package]]
+name = "const-random-macro"
+version = "0.1.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9d7d6ab3c3a2282db210df5f02c4dab6e0a7057af0fb7ebd4070f30fe05c0ddb"
+dependencies = [
+ "getrandom",
+ "once_cell",
+ "proc-macro-hack",
+ "tiny-keccak",
+]
+
 [[package]]
 name = "criterion"
 version = "0.3.6"
@@ -311,6 +369,12 @@ dependencies = [
  "cfg-if",
 ]
 
+[[package]]
+name = "crunchy"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
+
 [[package]]
 name = "csv"
 version = "1.1.6"
@@ -413,6 +477,12 @@ dependencies = [
  "termcolor",
 ]
 
+[[package]]
+name = "fixedbitset"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
+
 [[package]]
 name = "geigen"
 version = "0.1.0"
@@ -490,6 +560,15 @@ dependencies = [
  "thread_local",
 ]
 
+[[package]]
+name = "index_vec"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74086667896a940438f2118212f313abba4aff3831fef6f4b17d02add5c8bb60"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "indexmap"
 version = "1.9.2"
@@ -498,6 +577,7 @@ checksum = "1885e79c1fc4b10f0e172c475f458b7f7b93061064d98c3293e98c5ba0c8b399"
 dependencies = [
  "autocfg",
  "hashbrown",
+ "serde",
 ]
 
 [[package]]
@@ -677,6 +757,7 @@ dependencies = [
  "num-traits",
  "rawpointer",
  "rayon",
+ "serde",
 ]
 
 [[package]]
@@ -870,6 +951,18 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099"
 
+[[package]]
+name = "petgraph"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4dd7d28ee937e54fe3080c91faa1c3a46c06de6252988a7f4592ba2310ef22a4"
+dependencies = [
+ "fixedbitset",
+ "indexmap",
+ "serde",
+ "serde_derive",
+]
+
 [[package]]
 name = "plotters"
 version = "0.3.4"
@@ -919,6 +1012,12 @@ dependencies = [
  "num-integer",
 ]
 
+[[package]]
+name = "proc-macro-hack"
+version = "0.5.20+deprecated"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc375e1527247fe1a97d8b7156678dfe7c1af2fc075c9a4db3690ecd2a148068"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.51"
@@ -1167,12 +1266,16 @@ name = "scalib"
 version = "0.1.0"
 dependencies = [
  "approx",
+ "ariadne",
  "blis-src",
  "blis-sys",
+ "chumsky",
  "criterion",
  "crossbeam-utils",
  "geigen",
  "hytra",
+ "index_vec",
+ "indexmap",
  "indicatif",
  "itertools",
  "matrixmultiply",
@@ -1183,10 +1286,11 @@ dependencies = [
  "nshare",
  "num-integer",
  "num-traits",
- "numpy",
+ "petgraph",
  "pyo3",
  "rayon",
  "realfft 2.0.1",
+ "serde",
  "thiserror",
 ]
 
@@ -1194,6 +1298,7 @@ dependencies = [
 name = "scalib-py"
 version = "0.1.0"
 dependencies = [
+ "bincode",
  "indicatif",
  "ndarray",
  "num_cpus",
@@ -1221,6 +1326,9 @@ name = "serde"
 version = "1.0.152"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bb7d1f0d3021d347a83e556fc4683dea2ea09d87bccdf88ff5c12545d89d5efb"
+dependencies = [
+ "serde_derive",
+]
 
 [[package]]
 name = "serde_cbor"
@@ -1354,6 +1462,15 @@ dependencies = [
  "once_cell",
 ]
 
+[[package]]
+name = "tiny-keccak"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
+dependencies = [
+ "crunchy",
+]
+
 [[package]]
 name = "tinytemplate"
 version = "1.2.1"
@@ -1622,3 +1739,9 @@ name = "windows_x86_64_msvc"
 version = "0.42.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "447660ad36a13288b1db4d4248e857b510e8c3a225c822ba4fb748c0aafecffd"
+
+[[package]]
+name = "yansi"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09041cd90cf85f7f8b2df60c646f853b7f535ce68f85244eb6731cf89fa498ec"
diff --git a/src/scalib_ext/scalib-py/Cargo.toml b/src/scalib_ext/scalib-py/Cargo.toml
index d2f16265..c12a3683 100644
--- a/src/scalib_ext/scalib-py/Cargo.toml
+++ b/src/scalib_ext/scalib-py/Cargo.toml
@@ -22,6 +22,7 @@ ndarray = { version = "0.15.1", features = ["rayon"] }
 numpy = "0.18.0"
 indicatif = "0.17.3"
 num_cpus = "1.13.1"
+bincode = "1.3.3" # Serialization for pickle support
 
 [dependencies.pyo3]
 version = "0.18.0"
diff --git a/src/scalib_ext/scalib-py/src/belief_propagation.rs b/src/scalib_ext/scalib-py/src/belief_propagation.rs
index d0cbb1e4..c8c8a72c 100644
--- a/src/scalib_ext/scalib-py/src/belief_propagation.rs
+++ b/src/scalib_ext/scalib-py/src/belief_propagation.rs
@@ -50,32 +50,35 @@ pub fn to_var(function: &PyDict) -> Var {
 /// Convert the python description of a function node to a Func.
 pub fn to_func(function: &PyDict) -> Func {
     let neighboors: Vec<isize> = function.get_item("neighboors").unwrap().extract().unwrap();
-    let func: usize = function.get_item("func").unwrap().extract().unwrap();
+    let func: &str = function.get_item("func").unwrap().extract().unwrap();
 
     let neighboors: Vec<usize> = neighboors.iter().map(|x| *x as usize).collect();
 
     let f: FuncType;
-    if func == 0 {
+    if func == "AND" {
         f = FuncType::AND;
-    } else if func == 1 {
+    } else if func == "XOR" {
         f = FuncType::XOR;
-    } else if func == 2 {
+    } else if func == "NOT" {
+        f = FuncType::NOT;
+    } else if func == "XORCST" {
         let values: PyReadonlyArray1<u32> = function.get_item("values").unwrap().extract().unwrap();
         f = FuncType::XORCST(values.as_array().to_owned());
-    } else if func == 3 {
+    } else if func == "LOOKUP" {
         let table: PyReadonlyArray1<u32> = function.get_item("table").unwrap().extract().unwrap();
-        f = FuncType::LOOKUP(table.as_array().to_owned());
-    } else if func == 4 {
+        let table = table.as_array().to_owned();
+        f = FuncType::LOOKUP { table };
+    } else if func == "ANDCST" {
         let values: PyReadonlyArray1<u32> = function.get_item("values").unwrap().extract().unwrap();
         f = FuncType::ANDCST(values.as_array().to_owned());
-    } else if func == 5 {
+    } else if func == "ADD" {
         f = FuncType::ADD;
-    } else if func == 6 {
+    } else if func == "ADDCST" {
         let values: PyReadonlyArray1<u32> = function.get_item("values").unwrap().extract().unwrap();
         f = FuncType::ADDCST(values.as_array().to_owned());
-    } else if func == 7 {
+    } else if func == "MUL" {
         f = FuncType::MUL;
-    } else if func == 8 {
+    } else if func == "MULCST" {
         let values: PyReadonlyArray1<u32> = function.get_item("values").unwrap().extract().unwrap();
         f = FuncType::MULCST(values.as_array().to_owned());
     } else {
diff --git a/src/scalib_ext/scalib-py/src/factor_graph.rs b/src/scalib_ext/scalib-py/src/factor_graph.rs
new file mode 100644
index 00000000..ee8cf907
--- /dev/null
+++ b/src/scalib_ext/scalib-py/src/factor_graph.rs
@@ -0,0 +1,337 @@
+//! Python binding of SCALib's FactorGraph rust implementation.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use bincode::{deserialize, serialize};
+use numpy::{PyArray, PyArray1, PyArray2};
+use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
+use pyo3::prelude::*;
+use pyo3::types::{PyBytes, PyTuple};
+
+use scalib::sasca;
+
+#[pyclass(module = "_scalib_ext")]
+pub(crate) struct FactorGraph {
+    inner: Option<Arc<sasca::FactorGraph>>,
+}
+impl FactorGraph {
+    fn get_inner(&self) -> &Arc<sasca::FactorGraph> {
+        self.inner.as_ref().unwrap()
+    }
+    fn get_var(&self, var: &str) -> PyResult<sasca::VarId> {
+        self.get_inner()
+            .get_varid(var)
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+    fn get_factor(&self, factor: &str) -> PyResult<sasca::FactorId> {
+        self.get_inner()
+            .get_factorid(factor)
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+}
+
+// TODO run stuff on SCALib thread pool
+
+#[pymethods]
+impl FactorGraph {
+    #[new]
+    #[pyo3(signature = (*args))]
+    fn new(args: &PyTuple) -> PyResult<Self> {
+        if args.len() == 0 {
+            Ok(Self { inner: None })
+        } else {
+            let (description, tables): (
+                &str,
+                std::collections::HashMap<String, &PyArray1<sasca::ClassVal>>,
+            ) = args.extract()?;
+            let tables = tables
+                .into_iter()
+                .map(|(k, v)| PyResult::<_>::Ok((k, PyArray::to_vec(v)?)))
+                .collect::<Result<std::collections::HashMap<_, _>, _>>()?;
+            let fg = sasca::build_graph(description, tables)
+                .map_err(|e| PyValueError::new_err(e.to_string()))?;
+            Ok(Self {
+                inner: Some(Arc::new(fg)),
+            })
+        }
+    }
+
+    pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
+        Ok(PyBytes::new(py, &serialize(&self.inner.as_deref()).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()));
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn new_bp(&self, py: Python, nmulti: u32, public_values: PyObject) -> PyResult<BPState> {
+        let pub_values = pyobj2pubs(py, public_values, self.get_inner().public_multi())?;
+        Ok(BPState {
+            inner: Some(sasca::BPState::new(
+                self.get_inner().clone(),
+                nmulti,
+                pub_values,
+            )),
+        })
+    }
+
+    pub fn var_names(&self) -> Vec<&str> {
+        self.get_inner().var_names().collect()
+    }
+    pub fn factor_names(&self) -> Vec<&str> {
+        self.get_inner().factor_names().collect()
+    }
+    pub fn factor_scope<'s>(&'s self, factor: &str) -> PyResult<Vec<&'s str>> {
+        let factor_id = self.get_factor(factor)?;
+        Ok(self
+            .get_inner()
+            .factor_scope(factor_id)
+            .map(|v| self.get_inner().var_name(v))
+            .collect())
+    }
+    pub fn sanity_check(
+        &self,
+        py: Python,
+        public_values: PyObject,
+        var_assignments: PyObject,
+    ) -> PyResult<()> {
+        let inner = self.get_inner();
+        let pub_values = pyobj2pubs(py, public_values, inner.public_multi())?;
+        let var_values = pyobj2pubs(
+            py,
+            var_assignments,
+            inner.vars().map(|(v, vn)| (vn, inner.var_multi(v))),
+        )?;
+        inner
+            .sanity_check(pub_values, var_values.into())
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+}
+
+fn pyobj2pubs<'a>(
+    py: Python,
+    public_values: PyObject,
+    expected: impl Iterator<Item = (&'a str, bool)>,
+) -> PyResult<Vec<sasca::PublicValue>> {
+    let mut public_values: HashMap<&str, PyObject> = public_values.extract(py)?;
+    let pubs = expected
+        .map(|(pub_name, multi)| {
+            obj2pub(
+                py,
+                public_values
+                    .remove(pub_name)
+                    .ok_or_else(|| PyKeyError::new_err(format!("Missing value {}.", pub_name)))?,
+                multi,
+            )
+        })
+        .collect::<Result<Vec<sasca::PublicValue>, PyErr>>()?;
+    if public_values.is_empty() {
+        Ok(pubs)
+    } else {
+        let unknown_pubs = public_values.keys().collect::<Vec<_>>();
+        Err(PyKeyError::new_err(if unknown_pubs.len() == 1 {
+            format!("{} is not a public.", unknown_pubs[0])
+        } else {
+            format!("{:?} are not publics.", unknown_pubs)
+        }))
+    }
+}
+
+#[pyclass(module = "_scalib_ext")]
+pub(crate) struct BPState {
+    inner: Option<sasca::BPState>,
+}
+impl BPState {
+    fn get_inner(&self) -> &sasca::BPState {
+        self.inner.as_ref().unwrap()
+    }
+    fn get_inner_mut(&mut self) -> &mut sasca::BPState {
+        self.inner.as_mut().unwrap()
+    }
+    fn get_var(&self, var: &str) -> PyResult<sasca::VarId> {
+        self.get_inner()
+            .get_graph()
+            .get_varid(var)
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+    fn get_factor(&self, factor: &str) -> PyResult<sasca::FactorId> {
+        self.get_inner()
+            .get_graph()
+            .get_factorid(factor)
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+    fn get_edge(&self, var: sasca::VarId, factor: sasca::FactorId) -> PyResult<sasca::EdgeId> {
+        self.get_inner()
+            .get_graph()
+            .edge(var, factor)
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+    fn get_edge_named(&self, var: &str, factor: &str) -> PyResult<sasca::EdgeId> {
+        self.get_edge(self.get_var(var)?, self.get_factor(factor)?)
+    }
+}
+
+#[pymethods]
+impl BPState {
+    #[new]
+    #[pyo3(signature = (*_args))]
+    fn new(_args: &PyTuple) -> PyResult<Self> {
+        Ok(Self { inner: None })
+    }
+
+    pub fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
+        Ok(PyBytes::new(py, &serialize(&self.inner).unwrap()).to_object(py))
+    }
+
+    pub fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
+        match state.extract::<&PyBytes>(py) {
+            Ok(s) => {
+                self.inner = deserialize(s.as_bytes()).unwrap();
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn is_cyclic(&self) -> bool {
+        self.get_inner().is_cyclic()
+    }
+
+    pub fn set_evidence(&mut self, py: Python, var: &str, distr: PyObject) -> PyResult<()> {
+        let var_id = self.get_var(var)?;
+        let bp = self.get_inner_mut();
+        let distr = obj2distr(py, distr, bp.get_graph().var_multi(var_id))?;
+        bp.set_evidence(var_id, distr)
+            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
+        Ok(())
+    }
+    pub fn drop_evidence(&mut self, var: &str) -> PyResult<()> {
+        let var_id = self.get_var(var)?;
+        self.get_inner_mut().drop_evidence(var_id);
+        Ok(())
+    }
+    pub fn get_state(&self, py: Python, var: &str) -> PyResult<PyObject> {
+        distr2py(py, self.get_inner().get_state(self.get_var(var)?))
+    }
+    pub fn set_state(&mut self, py: Python, var: &str, distr: PyObject) -> PyResult<()> {
+        let var_id = self.get_var(var)?;
+        let bp = self.get_inner_mut();
+        let distr = obj2distr(py, distr, bp.get_graph().var_multi(var_id))?;
+        bp.set_state(var_id, distr)
+            .map_err(|e| PyTypeError::new_err(e.to_string()))?;
+        Ok(())
+    }
+    pub fn drop_state(&mut self, var: &str) -> PyResult<()> {
+        let var_id = self.get_var(var)?;
+        self.get_inner_mut().drop_state(var_id);
+        Ok(())
+    }
+    pub fn get_belief_to_var(&self, py: Python, var: &str, factor: &str) -> PyResult<PyObject> {
+        let edge_id = self.get_edge_named(var, factor)?;
+        distr2py(py, self.get_inner().get_belief_to_var(edge_id))
+    }
+    pub fn get_belief_from_var(&self, py: Python, var: &str, factor: &str) -> PyResult<PyObject> {
+        let edge_id = self.get_edge_named(var, factor)?;
+        distr2py(py, self.get_inner().get_belief_from_var(edge_id))
+    }
+    pub fn propagate_var(&mut self, var: &str) -> PyResult<()> {
+        let var_id = self.get_var(var)?;
+        self.get_inner_mut().propagate_var(var_id);
+        Ok(())
+    }
+    pub fn propagate_all_vars(&mut self) -> PyResult<()> {
+        self.get_inner_mut().propagate_all_vars();
+        Ok(())
+    }
+    pub fn propagate_factor_all(&mut self, factor: &str) -> PyResult<()> {
+        let factor_id = self.get_factor(factor)?;
+        self.get_inner_mut().propagate_factor_all(factor_id);
+        Ok(())
+    }
+    pub fn propagate_factor(
+        &mut self,
+        factor: &str,
+        dest: Vec<&str>,
+        clear_incoming: bool,
+    ) -> PyResult<()> {
+        let factor_id = self.get_factor(factor)?;
+        let dest = dest
+            .iter()
+            .map(|v| self.get_var(v))
+            .collect::<Result<Vec<_>, _>>()?;
+        self.get_inner_mut()
+            .propagate_factor(factor_id, dest.as_slice(), clear_incoming);
+        Ok(())
+    }
+    pub fn propagate_loopy_step(&mut self, n_steps: u32) {
+        self.get_inner_mut().propagate_loopy_step(n_steps);
+    }
+    pub fn graph(&self) -> FactorGraph {
+        FactorGraph {
+            inner: Some(self.get_inner().get_graph().clone()),
+        }
+    }
+    pub fn propagate_acyclic(
+        &mut self,
+        dest: &str,
+        clear_intermediates: bool,
+        clear_evidence: bool,
+    ) -> PyResult<()> {
+        let var = self.get_var(dest)?;
+        self.get_inner_mut()
+            .propagate_acyclic(var, clear_intermediates, clear_evidence)
+            .map_err(|e| PyValueError::new_err(e.to_string()))
+    }
+}
+
+fn obj2distr(py: Python, distr: PyObject, multi: bool) -> PyResult<sasca::Distribution> {
+    if multi {
+        let distr: &PyArray2<f64> = distr.extract(py)?;
+        sasca::Distribution::from_array_multi(
+            distr
+                .readonly()
+                .as_array()
+                .as_standard_layout()
+                .into_owned(),
+        )
+    } else {
+        let distr: &PyArray1<f64> = distr.extract(py)?;
+        sasca::Distribution::from_array_single(
+            distr
+                .readonly()
+                .as_array()
+                .as_standard_layout()
+                .into_owned(),
+        )
+    }
+    .map_err(|e| PyTypeError::new_err(e.to_string()))
+}
+
+fn obj2pub(py: Python, obj: PyObject, multi: bool) -> PyResult<sasca::PublicValue> {
+    Ok(if multi {
+        let obj: Vec<sasca::ClassVal> = obj.extract(py)?;
+        sasca::PublicValue::Multi(obj)
+    } else {
+        let obj: sasca::ClassVal = obj.extract(py)?;
+        sasca::PublicValue::Single(obj)
+    })
+}
+
+fn distr2py(py: Python, distr: &sasca::Distribution) -> PyResult<PyObject> {
+    if let Some(d) = distr.value() {
+        if distr.multi() {
+            return Ok(PyArray2::from_array(py, &d).into_py(py));
+        } else {
+            return Ok(PyArray1::from_array(py, &d.slice(ndarray::s![0, ..])).into_py(py));
+        }
+    } else {
+        return Ok(py.None());
+    }
+}
diff --git a/src/scalib_ext/scalib-py/src/lib.rs b/src/scalib_ext/scalib-py/src/lib.rs
index c1b60e25..6a9c620f 100644
--- a/src/scalib_ext/scalib-py/src/lib.rs
+++ b/src/scalib_ext/scalib-py/src/lib.rs
@@ -7,6 +7,7 @@ use pyo3::prelude::*;
 use std::error::Error;
 
 mod belief_propagation;
+mod factor_graph;
 mod lda;
 mod ranking;
 mod snr;
@@ -112,6 +113,8 @@ fn _scalib_ext(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<ttest::MTtest>()?;
     m.add_class::<lda::LDA>()?;
     m.add_class::<lda::LdaAcc>()?;
+    m.add_class::<factor_graph::FactorGraph>()?;
+    m.add_class::<factor_graph::BPState>()?;
     m.add_class::<thread_pool::ThreadPool>()?;
     m.add_function(wrap_pyfunction!(ranking::rank_accuracy, m)?)?;
     m.add_function(wrap_pyfunction!(ranking::rank_nbin, m)?)?;
diff --git a/src/scalib_ext/scalib/Cargo.toml b/src/scalib_ext/scalib/Cargo.toml
index fa5b91a5..db06e82a 100644
--- a/src/scalib_ext/scalib/Cargo.toml
+++ b/src/scalib_ext/scalib/Cargo.toml
@@ -1,8 +1,11 @@
 [package]
 name = "scalib"
 version = "0.1.0"
-authors = ["obronchain <olivier.bronchain@uclouvain.be>"]
-edition = "2018"
+authors = [
+    "Olivier Bronchain <olivier.bronchain@uclouvain.be>",
+    "Gaëtan Cassiers <gaetan.cassiers@uclouvain.be>",
+]
+edition = "2021"
 
 [dependencies.pyo3]
 version = "0.18.0"
@@ -12,24 +15,41 @@ features = ["extension-module"]
 blis = ["blis-src", "blis-sys"]
 
 [dependencies]
-numpy = "0.18.0"
+
+# Utilities
+thiserror = "1.0"
+itertools = "0.10.0"
 rayon = "1.5.0"
 num-integer = "0.1.42"
+index_vec = { version = "0.1.3", features = ["serde"] } # Type-safe indices in FactorGraph
+indexmap = { version = "1.9.1", features = ["serde"] } # Factor graph data
+serde = { version = "1.0", features = ["derive", "rc"] } # Serialization for pickle support
+
+# Factor graph parsing and error display
+chumsky = "0.8.0"
+ariadne = "0.1.5"
+
+# analysis of factor graphs
+petgraph = {version = "0.6.2", default-features=false, features = ["serde-1"]}
+
+
+# Progress bars
 indicatif = { version = "0.17.3" }
-ndarray = { version = "0.15.1", features = ["rayon", "approx"] }
-itertools = "0.10.0"
-ndarray-stats = "0.5.0"
-nalgebra = "0.26.1"
+hytra = "0.1.2"
+crossbeam-utils = "0.8.3"
+
 geigen = { path = "../geigen" }
-nshare = { git= "https://github.com/cassiersg/nshare", branch="rm_lt", features=["nalgebra", "ndarray", "std"] }
+
+# Numerical libs
+ndarray = { version = "0.15.1", features = ["rayon", "approx", "serde"] }
 blis-src = { version = "0.2.1-alpha.0", default-features=false, features=["static", "pthreads"], optional = true }
 blis-sys = { git="https://github.com/cassiersg/blis-sys", optional = true }
 matrixmultiply = "0.3.1"
-hytra = "0.1.2"
-crossbeam-utils = "0.8.3"
+nalgebra = "0.26.1"
+nshare = { git= "https://github.com/cassiersg/nshare", branch="rm_lt", features=["nalgebra", "ndarray", "std"] }
+
 realfft = "2.0.1"
 num-traits = "0.2.15"
-thiserror = "1.0"
 
 [lib]
 bench = false
@@ -38,6 +58,7 @@ bench = false
 criterion = "0.3.4"
 ndarray-rand = "0.14.0"
 approx = "0.4.0"
+ndarray-stats = "0.5.0"
 
 [[bench]]
 name = "belief_propagation"
diff --git a/src/scalib_ext/scalib/src/belief_propagation.rs b/src/scalib_ext/scalib/src/belief_propagation.rs
index 394c834f..49169347 100644
--- a/src/scalib_ext/scalib/src/belief_propagation.rs
+++ b/src/scalib_ext/scalib/src/belief_propagation.rs
@@ -8,7 +8,7 @@
 //!
 //! The values on the factor graph are probability distribution of values in GF(2)^n.
 
-use ndarray::{s, Array1, Array2, Axis};
+use ndarray::{s, Array1, Array2, ArrayView1, ArrayViewMut1, Axis};
 use rayon::prelude::*;
 use realfft::num_complex::Complex;
 use realfft::RealFftPlanner;
@@ -52,6 +52,8 @@ pub enum FuncType {
     AND,
     /// Bitwise XOR of variables
     XOR,
+    /// Bitwise NOT of variables
+    NOT,
     /// Modular ADD of variables
     ADD,
     /// Modular MUL of variables
@@ -65,7 +67,7 @@ pub enum FuncType {
     /// Modular MUL of variables, MULing additionally a public variable.
     MULCST(Array1<u32>),
     /// Lookup table function.
-    LOOKUP(Array1<u32>),
+    LOOKUP { table: Array1<u32> },
 }
 
 /// A function node in the graph.
@@ -76,7 +78,7 @@ pub struct Func {
 }
 
 /// The minimum non-zero probability (to avoid denormalization, etc.)
-const MIN_PROBA: f64 = 1e-20;
+const MIN_PROBA: f64 = 1e-40;
 
 /// Clip down to `MIN_PROBA`
 fn make_non_zero<S: ndarray::DataMut + ndarray::RawData<Elem = f64>, D: ndarray::Dimension>(
@@ -85,6 +87,21 @@ fn make_non_zero<S: ndarray::DataMut + ndarray::RawData<Elem = f64>, D: ndarray:
     x.mapv_inplace(|y| y.max(MIN_PROBA));
 }
 
+fn make_non_zero_signed<
+    S: ndarray::DataMut + ndarray::RawData<Elem = f64>,
+    D: ndarray::Dimension,
+>(
+    x: &mut ndarray::ArrayBase<S, D>,
+) {
+    x.mapv_inplace(|y| {
+        if y.is_sign_positive() {
+            y.max(MIN_PROBA)
+        } else {
+            y.min(-MIN_PROBA)
+        }
+    });
+}
+
 /// Walsh-Hadamard transform (non-normalized).
 #[inline(always)]
 fn fwht(a: &mut [f64], len: usize) {
@@ -106,6 +123,69 @@ fn fwht(a: &mut [f64], len: usize) {
     }
 }
 
+/// Cumulative transform (U in RLDA paper)
+#[inline(always)]
+fn cumt(a: &mut [f64], len: usize) {
+    // Note: the speed of this can probably be much improved, with the following techiques
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x + y;
+                a[j + h] = y;
+            }
+        }
+        h *= 2;
+    }
+}
+
+/// Cumulative inverse transform (U^-1 in RLDA paper)
+#[inline(always)]
+fn cumti(a: &mut [f64], len: usize) {
+    // Note: the speed of this can probably be much improved, with the following techiques
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x - y;
+                a[j + h] = y;
+            }
+        }
+        h *= 2;
+    }
+}
+
+/// Tansform for operand of AND (V in RLDA paper), involutive
+#[inline(always)]
+fn opandt(a: &mut [f64], len: usize) {
+    // Note: the speed of this can probably be much improved, with the following techiques
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x;
+                a[j + h] = x - y;
+            }
+        }
+        h *= 2;
+    }
+}
+
 /// Make it such that the sum of the probabilities in the distribution is 1.0.
 /// `distri` can be a ParaDistri or a SingleDistri.
 fn normalize_distri(distri: &mut Array2<f64>) {
@@ -200,7 +280,8 @@ pub fn update_functions(functions: &[Func], edges: &mut [Vec<&mut Array2<f64>>])
         .for_each(|(function, edge)| match &function.functype {
             // TODO: if nc is prime, the update for MUL can be computed more efficiently by mapping
             // classes to their discrete logarithm, and by applying FFT.
-            FuncType::AND | FuncType::MUL => {
+            // If not prime, use CRT.
+            FuncType::MUL => {
                 let [output_msg, input1_msg, input2_msg]: &mut [_; 3] =
                     edge.as_mut_slice().try_into().unwrap();
                 let nc = input1_msg.shape()[1];
@@ -223,13 +304,7 @@ pub fn update_functions(functions: &[Func], edges: &mut [Vec<&mut Array2<f64>>])
                             for i1 in 0..nc {
                                 for i2 in 0..nc {
                                     // Unifies operators that can only be binary
-                                    let o = match &function.functype {
-                                        FuncType::AND => i1 & i2,
-                                        FuncType::MUL => {
-                                            (((i1 * i2) as u32) % (nc as u32)) as usize
-                                        }
-                                        _ => unreachable!(),
-                                    };
+                                    let o = (((i1 * i2) as u32) % (nc as u32)) as usize;
                                     in1_msg_scratch[i1] += input2_msg[i2] * output_msg[o];
                                     in2_msg_scratch[i2] += input1_msg[i1] * output_msg[o];
                                     out_msg_scratch[o] += input1_msg[i1] * input2_msg[i2];
@@ -241,12 +316,33 @@ pub fn update_functions(functions: &[Func], edges: &mut [Vec<&mut Array2<f64>>])
                         },
                     );
             }
+            FuncType::AND => {
+                ands(edge.as_mut());
+            }
             FuncType::ADD => {
                 adds(edge.as_mut());
             }
             FuncType::XOR => {
                 xors(edge.as_mut());
             }
+            FuncType::NOT => {
+                let [output_msg, input_msg]: &mut [_; 2] = edge.as_mut_slice().try_into().unwrap();
+                let nc = input_msg.shape()[1];
+                (input_msg.outer_iter_mut(), output_msg.outer_iter_mut())
+                    .into_par_iter()
+                    .for_each_init(
+                        || (Array1::zeros(nc), Array1::zeros(nc)),
+                        |(in_msg_scratch, out_msg_scratch), (mut input_msg, mut output_msg)| {
+                            for i in 0..nc {
+                                let o = (nc - 1) ^ i as usize;
+                                in_msg_scratch[i] = output_msg[o];
+                                out_msg_scratch[o] = input_msg[i];
+                            }
+                            input_msg.assign(in_msg_scratch);
+                            output_msg.assign(out_msg_scratch);
+                        },
+                    );
+            }
             FuncType::XORCST(values)
             | FuncType::ANDCST(values)
             | FuncType::ADDCST(values)
@@ -286,7 +382,7 @@ pub fn update_functions(functions: &[Func], edges: &mut [Vec<&mut Array2<f64>>])
                         },
                     );
             }
-            FuncType::LOOKUP(table) => {
+            FuncType::LOOKUP { table } => {
                 let [output_msg, input1_msg]: &mut [_; 2] = edge.as_mut_slice().try_into().unwrap();
                 let nc = input1_msg.shape()[1];
                 (input1_msg.outer_iter_mut(), output_msg.outer_iter_mut())
@@ -298,9 +394,6 @@ pub fn update_functions(functions: &[Func], edges: &mut [Vec<&mut Array2<f64>>])
                             out_msg_scratch.fill(0.0);
                             for i1 in 0..nc {
                                 let o: usize = table[i1] as usize;
-                                // This requires table to be bijective. Otherwise, we would have to
-                                // divide the messge on the output by the number of matching inputs
-                                // to get the message to forward on the input edge.
                                 in1_msg_scratch[i1] += output_msg[o];
                                 out_msg_scratch[o] += input1_msg[i1];
                             }
@@ -400,6 +493,59 @@ pub fn xors(inputs: &mut [&mut Array2<f64>]) {
     }
 }
 
+/// Compute a AND function node between all edges.
+pub fn ands(inputs: &mut [&mut Array2<f64>]) {
+    let n_runs = inputs[0].shape()[0];
+    let nc = inputs[0].shape()[1];
+    let (input_0, input_1, input_2) = match inputs {
+        [i0, i1, i2] => (i0, i1, i2),
+        _ => panic!("Current implementation limit AND to 2 inputs."),
+    };
+    for run in 0..n_runs {
+        let mut cum_t1 = input_1.slice(s![run, ..]).to_owned();
+        let mut cum_t2 = input_2.slice(s![run, ..]).to_owned();
+        let cum_transform = |mut x: ArrayViewMut1<f64>| {
+            let x_slice = x.as_slice_mut().unwrap();
+            cumt(x_slice, nc);
+            make_non_zero(&mut x);
+        };
+        let opand_transform = |mut x: ArrayViewMut1<f64>| {
+            let x_slice = x.as_slice_mut().unwrap();
+            opandt(x_slice, nc);
+            make_non_zero_signed(&mut x);
+        };
+        cum_transform(cum_t1.view_mut());
+        cum_transform(cum_t2.view_mut());
+        opand_transform(input_0.slice_mut(s![run, ..]));
+        let icumt_norm = |mut x: ArrayViewMut1<f64>| {
+            let x_slice = x.as_slice_mut().unwrap();
+            cumti(x_slice, nc);
+            let s = 1e-100f64.max(x.sum());
+            x /= s;
+        };
+        let v0 = input_0.slice(s![run, ..]);
+        let set_input_distr = |mut op: ArrayViewMut1<f64>, t_other: ArrayView1<f64>| {
+            ndarray::Zip::from(op.view_mut())
+                .and(v0)
+                .and(t_other)
+                .for_each(|input, res, other| *input = *res * *other);
+            make_non_zero_signed(&mut op);
+            let op_slice = op.as_slice_mut().unwrap();
+            opandt(op_slice, nc);
+            let s = op.sum();
+            op /= s;
+            make_non_zero_signed(&mut op);
+        };
+        set_input_distr(input_1.slice_mut(s![run, ..]), cum_t2.view());
+        set_input_distr(input_2.slice_mut(s![run, ..]), cum_t1.view());
+        ndarray::Zip::from(input_0.slice_mut(s![run, ..]))
+            .and(cum_t1.view())
+            .and(cum_t2.view())
+            .for_each(|res, in1, in2| *res = *in1 * *in2);
+        icumt_norm(input_0.slice_mut(s![run, ..]));
+    }
+}
+
 /// Run the belief propagation algorithm on the python representation of a factor graph.
 pub fn run_bp(
     functions: &[Func],
diff --git a/src/scalib_ext/scalib/src/lda.rs b/src/scalib_ext/scalib/src/lda.rs
index 37ac6a34..b66ff18c 100644
--- a/src/scalib_ext/scalib/src/lda.rs
+++ b/src/scalib_ext/scalib/src/lda.rs
@@ -57,7 +57,12 @@ impl LdaAcc {
     }
 
     /// Traces: shape (n, ns). Classes shape: (n,)
-    fn new(nc: usize, traces: ArrayView2<i16>, classes: ArrayView1<u16>, gemm_algo: u32) -> Self {
+    pub fn new(
+        nc: usize,
+        traces: ArrayView2<i16>,
+        classes: ArrayView1<u16>,
+        gemm_algo: u32,
+    ) -> Self {
         let mut res = Self::from_dim(nc, traces.shape()[1]);
         res.update(traces, classes, gemm_algo);
         return res;
diff --git a/src/scalib_ext/scalib/src/lib.rs b/src/scalib_ext/scalib/src/lib.rs
index a02252bb..6e05803f 100644
--- a/src/scalib_ext/scalib/src/lib.rs
+++ b/src/scalib_ext/scalib/src/lib.rs
@@ -2,6 +2,7 @@ pub mod belief_propagation;
 pub mod lda;
 pub(crate) mod matrixmul;
 pub mod mttest;
+pub mod sasca;
 pub mod snr;
 pub mod ttest;
 pub(crate) mod utils;
diff --git a/src/scalib_ext/scalib/src/sasca/belief_propagation.rs b/src/scalib_ext/scalib/src/sasca/belief_propagation.rs
new file mode 100644
index 00000000..65b41322
--- /dev/null
+++ b/src/scalib_ext/scalib/src/sasca/belief_propagation.rs
@@ -0,0 +1,718 @@
+use std::sync::Arc;
+
+use thiserror::Error;
+
+use super::factor_graph::{
+    EdgeId, EdgeSlice, EdgeVec, Factor, FactorId, FactorKind, FactorVec, Node, Table, VarId, VarVec,
+};
+use super::{Distribution, FactorGraph, PublicValue};
+
+// TODO improvements
+// - use a pool for Distribution allocations (can be a simple Vec storing them), to avoid frequent
+// allocations
+
+// Workaround since the plans are not Serialize of Debug
+#[derive(Clone, serde::Serialize, serde::Deserialize)]
+#[serde(from = "FftPlansSer", into = "FftPlansSer")]
+pub(super) struct FftPlans {
+    pub(super) size: usize,
+    pub(super) r2c: Arc<dyn realfft::RealToComplex<f64>>,
+    pub(super) c2r: Arc<dyn realfft::ComplexToReal<f64>>,
+}
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+struct FftPlansSer {
+    size: usize,
+}
+
+impl From<FftPlansSer> for FftPlans {
+    fn from(p: FftPlansSer) -> Self {
+        Self::new(p.size)
+    }
+}
+
+impl From<FftPlans> for FftPlansSer {
+    fn from(p: FftPlans) -> Self {
+        Self { size: p.size }
+    }
+}
+
+impl std::fmt::Debug for FftPlans {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+        f.debug_struct("FftPlans")
+            .field("size", &self.size)
+            .finish_non_exhaustive()
+    }
+}
+
+impl FftPlans {
+    fn new(size: usize) -> Self {
+        let mut planner = realfft::RealFftPlanner::new();
+        Self {
+            size,
+            r2c: planner.plan_fft_forward(size),
+            c2r: planner.plan_fft_inverse(size),
+        }
+    }
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct BPState {
+    graph: std::sync::Arc<FactorGraph>,
+    // number of parallel executions for PARA vars
+    nmulti: u32,
+    // one public for every factor. Set to 0 if not relevant.
+    public_values: FactorVec<PublicValue>,
+    // public value for each factor
+    pub_reduced: FactorVec<PublicValue>,
+    // evidence for each var
+    evidence: VarVec<Distribution>,
+    // current proba for each var
+    var_state: VarVec<Distribution>,
+    // beliefs on each edge
+    belief_from_var: EdgeVec<Distribution>,
+    belief_to_var: EdgeVec<Distribution>,
+    // save if cyclic
+    cyclic: bool,
+    // fft plans
+    plans: FftPlans,
+}
+
+#[derive(Debug, Clone, Error)]
+pub enum BPError {
+    #[error("Wrong distribution kind: got {0}, expected {1}.")]
+    WrongDistributionKind(&'static str, &'static str),
+    #[error("Wrong number of classes for distribution: got {0}, expected {1}.")]
+    WrongDistributionNc(usize, usize),
+    #[error("Wrong number of traces for distribution: got {0}, expected {1}.")]
+    WrongDistributionNmulti(usize, u32),
+    #[error("The distribution is not is C memory order. Shape: {0:?}, strides: {1:?}.")]
+    DistributionLayout(Vec<usize>, Vec<isize>),
+    #[error("Cannot run acyclic BP on a cyclic graph.")]
+    NotAcyclic,
+}
+
+impl BPState {
+    pub fn new(
+        graph: std::sync::Arc<FactorGraph>,
+        nmulti: u32,
+        public_values: Vec<PublicValue>,
+    ) -> Self {
+        let var_state: VarVec<_> = graph
+            .vars
+            .values()
+            .map(|v| Distribution::new(v.multi, graph.nc, nmulti))
+            .collect();
+        let beliefs: EdgeVec<_> = graph
+            .edges
+            .iter()
+            .map(|e| Distribution::new(graph.factor(e.factor).multi, graph.nc, nmulti))
+            .collect();
+        let pub_reduced = graph.reduce_pub(&public_values);
+        let cyclic = graph.is_cyclic(nmulti > 1);
+        let plans = FftPlans::new(graph.nc);
+        Self {
+            evidence: var_state.clone(),
+            belief_from_var: beliefs.clone(),
+            belief_to_var: beliefs,
+            var_state,
+            graph,
+            nmulti,
+            public_values: FactorVec::from_vec(public_values),
+            pub_reduced,
+            cyclic,
+            plans,
+        }
+    }
+    pub fn is_cyclic(&self) -> bool {
+        self.cyclic
+    }
+    pub fn get_graph(&self) -> &std::sync::Arc<FactorGraph> {
+        &self.graph
+    }
+    fn check_distribution(&self, distr: &Distribution, multi: bool) -> Result<(), BPError> {
+        if distr.multi() != multi {
+            Err(BPError::WrongDistributionKind(
+                if distr.multi() { "multi" } else { "single" },
+                if multi { "multi" } else { "single" },
+            ))
+        } else if distr.shape().1 != self.graph.nc {
+            Err(BPError::WrongDistributionNc(distr.shape().1, self.graph.nc))
+        } else if distr.multi() && self.nmulti as usize != distr.shape().0 {
+            Err(BPError::WrongDistributionNmulti(
+                distr.shape().0,
+                self.nmulti,
+            ))
+        } else {
+            Ok(())
+        }
+    }
+    pub fn set_evidence(&mut self, var: VarId, evidence: Distribution) -> Result<(), BPError> {
+        self.check_distribution(&evidence, self.graph.var_multi(var))?;
+        self.evidence[var] = evidence;
+        Ok(())
+    }
+    pub fn drop_evidence(&mut self, var: VarId) {
+        self.evidence[var] = self.evidence[var].as_uniform();
+    }
+    pub fn get_state(&self, var: VarId) -> &Distribution {
+        &self.var_state[var]
+    }
+    pub fn set_state(&mut self, var: VarId, state: Distribution) -> Result<(), BPError> {
+        self.check_distribution(&state, self.graph.var_multi(var))?;
+        self.var_state[var] = state;
+        Ok(())
+    }
+    pub fn drop_state(&mut self, var: VarId) {
+        self.var_state[var] = self.var_state[var].as_uniform();
+    }
+    pub fn get_belief_to_var(&self, edge: EdgeId) -> &Distribution {
+        &self.belief_to_var[edge]
+    }
+    pub fn get_belief_from_var(&self, edge: EdgeId) -> &Distribution {
+        &self.belief_from_var[edge]
+    }
+    // Propagation type:
+    // belief to var -> var
+    // var -> belief to func
+    // trhough func: towards all vars, towards a subset of vars
+    pub fn propagate_to_var(&mut self, var: VarId, clear_evidence: bool) {
+        let distr_iter = self
+            .graph
+            .var(var)
+            .edges
+            .values()
+            .map(|e| &self.belief_to_var[*e]);
+        self.var_state[var] = self.evidence[var].take_or_clone(clear_evidence);
+        // We multiply_reg to avoid having very low values in the product.
+        // Since inputs should not be too big, we should not have any overflow.
+        // Underflow my happen, since probas are lower-bounded by MIN_PROBA**2.
+        self.var_state[var].multiply_reg(distr_iter);
+        // Now we'll make to sum equal one to avoid underflows or overflows in the long run, and keep the exposed probas nice.
+        // Just multiply, don't add anything, underflows are taken care of above.
+        self.var_state[var].normalize();
+    }
+    pub fn propagate_from_var(&mut self, edge: EdgeId) {
+        // Dividing here is ok if we ensure that there is no zero element and no
+        // underflow (or denormalization).
+        // This is guaranteed as long as min_proba > var_degree * MIN_POSITIVE
+        let var = self.graph.edges[edge].var;
+        self.belief_from_var[edge].reset();
+        self.belief_from_var[edge] =
+            Distribution::divide_reg(&self.var_state[var], &self.belief_to_var[edge]);
+    }
+    pub fn propagate_factor(&mut self, factor_id: FactorId, dest: &[VarId], clear_incoming: bool) {
+        let factor = self.graph.factor(factor_id);
+        // Pre-erase to have buffers available in cache allocator.
+        for d in dest {
+            self.belief_to_var[factor.edges[d]].reset();
+        }
+        // Use a macro to call very similar functions in match arms.
+        // Needed because of anonymous return types of these functions.
+        macro_rules! prop_factor {
+            ($f:ident, $($arg:expr),*) => {
+                {
+                    let it = $f(factor, &mut self.belief_from_var, dest, clear_incoming, $($arg,)*);
+                    for (distr, dest) in it.zip(dest.iter()) {
+                        self.belief_to_var[factor.edges[dest]] = distr;
+                    }
+                }
+            };
+        }
+        match &factor.kind {
+            FactorKind::AND { .. } => {
+                prop_factor!(factor_gen_and, &self.pub_reduced[factor_id])
+            }
+            FactorKind::XOR => prop_factor!(factor_xor, &self.pub_reduced[factor_id]),
+            FactorKind::NOT => prop_factor!(factor_not, (self.graph.nc - 1) as u32),
+            FactorKind::ADD => prop_factor!(factor_add, &self.pub_reduced[factor_id], &self.plans),
+            FactorKind::MUL => prop_factor!(factor_mul, &self.pub_reduced[factor_id]),
+            FactorKind::LOOKUP { table } => {
+                prop_factor!(factor_lookup, &self.graph.tables[*table])
+            }
+        }
+    }
+
+    // Higher-level
+    pub fn propagate_factor_all(&mut self, factor: FactorId) {
+        let dest: Vec<_> = self.graph.factor(factor).edges.keys().cloned().collect();
+        self.propagate_factor(factor, dest.as_slice(), false);
+    }
+    pub fn propagate_from_var_all(&mut self, var: VarId) {
+        for i in 0..self.graph.var(var).edges.len() {
+            self.propagate_from_var(self.graph.var(var).edges[i]);
+        }
+        for i in 0..self.graph.var(var).edges.len() {
+            self.belief_to_var[self.graph.var(var).edges[i]].reset();
+        }
+    }
+    pub fn propagate_var(&mut self, var: VarId) {
+        self.propagate_to_var(var, false);
+        self.propagate_from_var_all(var);
+    }
+    pub fn propagate_all_vars(&mut self) {
+        for var_id in self.graph.range_vars() {
+            self.propagate_var(var_id);
+        }
+    }
+    pub fn propagate_loopy_step(&mut self, n_steps: u32) {
+        for _ in 0..n_steps {
+            for factor_id in self.graph.range_factors() {
+                self.propagate_factor_all(factor_id);
+            }
+            self.propagate_all_vars();
+        }
+    }
+    pub fn propagate_acyclic(
+        &mut self,
+        var: VarId,
+        clear_intermediates: bool,
+        clear_evidence: bool,
+    ) -> Result<(), BPError> {
+        if self.is_cyclic() {
+            return Err(BPError::NotAcyclic);
+        }
+        for (node, parent) in self.graph.propagation_order(var) {
+            dbg!(&self);
+            dbg!(node);
+            dbg!(parent);
+            match node {
+                Node::Var(var_id) => {
+                    self.propagate_to_var(var_id, clear_evidence);
+                    if let Some(dest_factor) = parent {
+                        let edge_id = self.graph.var(var_id).edges[&dest_factor.factor().unwrap()];
+                        self.propagate_from_var(edge_id);
+                    }
+                    if clear_intermediates {
+                        for e in self.graph.var(var_id).edges.values() {
+                            self.belief_to_var[*e].reset();
+                        }
+                    }
+                }
+                Node::Factor(factor_id) => {
+                    let parent_var = parent.unwrap().var().unwrap();
+                    self.propagate_factor(factor_id, &[parent_var], clear_intermediates);
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+fn factor_gen_and<'a>(
+    factor: &'a Factor,
+    belief_from_var: &'a mut EdgeSlice<Distribution>,
+    dest: &'a [VarId],
+    clear_incoming: bool,
+    pub_red: &PublicValue,
+) -> impl Iterator<Item = Distribution> + 'a {
+    let FactorKind::AND { vars_neg } = &factor.kind else { unreachable!() };
+    // Special case for single-input AND
+    if factor.has_res & (factor.edges.len() == 2) {
+        return dest
+            .iter()
+            .map(|var| {
+                let i = factor.edges.get_index_of(var).unwrap();
+                let mut distr = belief_from_var[factor.edges[1 - i]].take_or_clone(clear_incoming);
+                if vars_neg[1 - i] {
+                    distr.not();
+                }
+                if i == 0 {
+                    // dest is the result of the AND
+                    distr.and_cst(pub_red);
+                } else {
+                    // dest is an operand of the AND, original distr is result
+                    distr.inv_and_cst(pub_red);
+                }
+                if vars_neg[i] {
+                    distr.not();
+                }
+                distr
+            })
+            .collect::<Vec<_>>()
+            .into_iter();
+    }
+    // Compute a product in the transformed domain
+    // direct transform:
+    // - if operand: cumt
+    // - if result: opandt
+    // inverse transform:
+    // - if operand: opandt^-1 = opandt
+    // - if result: cumt^-1 = cumti
+    let mut acc = belief_from_var[factor.edges[0]].new_constant(pub_red);
+    if factor.has_res {
+        // constant is operand
+        acc.cumt();
+    } else {
+        // constant is result
+        acc.opandt();
+    }
+    let mut taken_dest = vec![false; factor.edges.len()];
+    for dest in dest {
+        taken_dest[factor.edges.get_index_of(dest).unwrap()] = true;
+    }
+    // Compute the product in transform domain
+    // We do not take the product of all factors then divide because some factors could be zero.
+    let mut dest_transformed = Vec::with_capacity(dest.len());
+    for ((i, e), taken) in factor.edges.values().enumerate().zip(taken_dest.iter()) {
+        let mut d = belief_from_var[*e].take_or_clone(clear_incoming);
+        if vars_neg[i] {
+            d.not();
+        }
+        d.ensure_full();
+        if factor.has_res && (i == 0) {
+            d.opandt();
+        } else {
+            d.cumt();
+        }
+        // We either multiply (non-taken distributions) or we add to the vector of factors.
+        if !*taken {
+            acc.multiply(Some(&d).into_iter());
+        } else {
+            dest_transformed.push(d);
+        }
+    }
+    // This could be done in O(l log l) instead of O(l^2) where l=dest.len()
+    // by better caching product computations.
+    return (0..dest.len())
+        .map(|i| {
+            let mut res = acc.clone();
+            res.multiply(
+                (0..dest.len())
+                    .filter(|j| *j != i)
+                    .map(|j| &dest_transformed[j]),
+            );
+            // Inverse transform
+            if Some(dest[i]) == factor.res_id() {
+                res.cumti();
+            } else {
+                res.opandt();
+            }
+            if vars_neg[factor.edges.get_index_of(&dest[i]).unwrap()] {
+                res.not();
+            }
+            res.regularize();
+            res
+        })
+        .collect::<Vec<_>>()
+        .into_iter();
+}
+
+fn reset_incoming(
+    factor: &Factor,
+    belief_from_var: &mut EdgeSlice<Distribution>,
+    dest_taken: &[bool],
+    clear_incoming: bool,
+) {
+    // Everything will be uniform.
+    // Clear incoming and reset outgoing.
+    if clear_incoming {
+        for (taken, e) in dest_taken.iter().zip(factor.edges.values()) {
+            if *taken {
+                belief_from_var[*e].reset();
+            }
+        }
+    }
+}
+
+fn factor_xor<'a>(
+    factor: &'a Factor,
+    belief_from_var: &'a mut EdgeSlice<Distribution>,
+    dest: &'a [VarId],
+    clear_incoming: bool,
+    pub_red: &PublicValue,
+) -> impl Iterator<Item = Distribution> + 'a {
+    // Special case for single-input XOR
+    if factor.edges.len() == 2 {
+        return dest
+            .iter()
+            .map(|var| {
+                let i = factor.edges.get_index_of(var).unwrap();
+                let mut distr = belief_from_var[factor.edges[1 - i]].take_or_clone(clear_incoming);
+                distr.xor_cst(pub_red);
+                distr
+            })
+            .collect::<Vec<_>>()
+            .into_iter();
+    }
+    let mut acc = belief_from_var[factor.edges[0]].new_constant(pub_red);
+    acc.wht();
+    let mut taken_dest = vec![false; factor.edges.len()];
+    for dest in dest {
+        taken_dest[factor.edges.get_index_of(dest).unwrap()] = true;
+    }
+    let mut uniform_iter = factor
+        .edges
+        .values()
+        .zip(taken_dest.iter())
+        .enumerate()
+        .filter(|(_, (e, _))| !belief_from_var[**e].is_full());
+    let uniform_op = uniform_iter.next();
+    if let Some((i, (e_dest, t))) = uniform_op {
+        if !*t || uniform_iter.next().is_some() {
+            // At least 2 uniform operands, or single uniform is not in dest,
+            // all dest messages are uniform.
+            reset_incoming(factor, belief_from_var, &taken_dest, clear_incoming);
+            return vec![acc.as_uniform(); dest.len()].into_iter();
+        } else {
+            // Single uniform op, only compute for that one.
+            for e in factor.edges.values() {
+                if e != e_dest {
+                    let mut d = belief_from_var[*e].take_or_clone(clear_incoming);
+                    d.wht();
+                    d.make_non_zero_signed();
+                    acc.multiply(Some(&d).into_iter());
+                }
+            }
+            acc.wht();
+            acc.regularize();
+            let mut res = vec![acc.as_uniform(); dest.len()];
+            res[i] = acc;
+            return res.into_iter();
+        }
+    } else {
+        // Here we have to actually compute.
+        // Simply make the product if Walsh-Hadamard domain
+        // We do take the product of all factors then divide because some factors could be zero.
+        let mut dest_wht = Vec::with_capacity(dest.len());
+        for (e, taken) in factor.edges.values().zip(taken_dest.iter()) {
+            let mut d = belief_from_var[*e].take_or_clone(clear_incoming);
+            assert!(d.is_full());
+            d.wht();
+            // TODO remove this ?
+            d.make_non_zero_signed();
+            // We either multiply (non-taken distributions) or we add to the vector of factors.
+            if !*taken {
+                acc.multiply(Some(&d).into_iter());
+            } else {
+                dest_wht.push(d);
+            }
+        }
+        // This could be done in O(l log l) instead of O(l^2) where l=dest.len()
+        // by better caching product computations.
+        return (0..dest.len())
+            .map(|i| {
+                let mut res = acc.clone();
+                res.multiply((0..dest.len()).filter(|j| *j != i).map(|j| &dest_wht[j]));
+                res.wht();
+                res.regularize();
+                res
+            })
+            .collect::<Vec<_>>()
+            .into_iter();
+    }
+}
+
+fn factor_not<'a>(
+    factor: &'a Factor,
+    belief_from_var: &'a mut EdgeSlice<Distribution>,
+    dest: &'a [VarId],
+    clear_incoming: bool,
+    inv_cst: u32,
+) -> impl Iterator<Item = Distribution> + 'a {
+    factor_xor(
+        factor,
+        belief_from_var,
+        dest,
+        clear_incoming,
+        &PublicValue::Single(inv_cst),
+    )
+}
+
+// TODO handle subtraction too (actually, we can re-write it as an addition by moving terms around).
+fn factor_add<'a>(
+    factor: &'a Factor,
+    belief_from_var: &'a mut EdgeSlice<Distribution>,
+    dest: &'a [VarId],
+    clear_incoming: bool,
+    pub_red: &PublicValue,
+    plans: &FftPlans,
+) -> impl Iterator<Item = Distribution> + 'a {
+    // Special case for single-input ADD
+    if factor.edges.len() == 2 {
+        return dest
+            .iter()
+            .map(|var| {
+                let i = factor.edges.get_index_of(var).unwrap();
+                let mut distr = belief_from_var[factor.edges[1 - i]].take_or_clone(clear_incoming);
+                distr.add_cst(pub_red, i != 0);
+                distr
+            })
+            .collect::<Vec<_>>()
+            .into_iter();
+    }
+    let mut taken_dest = vec![false; factor.edges.len()];
+    for dest in dest {
+        taken_dest[factor.edges.get_index_of(dest).unwrap()] = true;
+    }
+    let mut uniform_iter = factor
+        .edges
+        .values()
+        .zip(taken_dest.iter())
+        .enumerate()
+        .filter(|(_, (e, _))| !belief_from_var[**e].is_full());
+    let uniform_op = uniform_iter.next();
+    let uniform_template = belief_from_var[factor.edges[0]].as_uniform();
+    let (nmulti, nc) = uniform_template.shape();
+    let mut fft_tmp = ndarray::Array2::zeros((nmulti, nc / 2 + 1));
+    let mut acc_fft = ndarray::Array2::zeros((nmulti, nc / 2 + 1));
+    let mut acc_fft_init = false;
+    let mut fft_scratch = plans.r2c.make_scratch_vec();
+    let mut fft_input_scratch = plans.r2c.make_input_vec();
+    if let Some((i, (e_dest, t))) = uniform_op {
+        if !*t || uniform_iter.next().is_some() {
+            // At least 2 uniform operands, or single uniform is not in dest,
+            // all dest messages are uniform.
+            reset_incoming(factor, belief_from_var, &taken_dest, clear_incoming);
+            return vec![uniform_template; dest.len()].into_iter();
+        } else {
+            // Single uniform op, only compute for that one.
+            for e in factor.edges.values() {
+                if e != e_dest {
+                    belief_from_var[*e].fft_to(
+                        fft_input_scratch.as_mut_slice(),
+                        fft_tmp.view_mut(),
+                        fft_scratch.as_mut_slice(),
+                        plans,
+                    );
+                    if acc_fft_init {
+                        acc_fft *= &fft_tmp;
+                    } else {
+                        acc_fft.assign(&fft_tmp);
+                        acc_fft_init = true;
+                    }
+                    if clear_incoming {
+                        belief_from_var[*e].reset();
+                    }
+                }
+            }
+            let mut acc = uniform_template.clone();
+            let mut fft_scratch = plans.c2r.make_scratch_vec();
+            acc.ifft(acc_fft.view_mut(), fft_scratch.as_mut_slice(), plans);
+            acc.regularize();
+            let mut res = vec![uniform_template; dest.len()];
+            res[i] = acc;
+            return res.into_iter();
+        }
+    } else {
+        // Here we have to actually compute.
+        // Simply make the product if FFT domain
+        // We do take the product of all factors then divide because some factors could be zero.
+        let mut dest_fft = Vec::with_capacity(dest.len());
+        for (e, taken) in factor.edges.values().zip(taken_dest.iter()) {
+            let v = if *taken {
+                let mut fft_e = ndarray::Array2::zeros((nmulti, nc / 2 + 1));
+                belief_from_var[*e].fft_to(
+                    fft_input_scratch.as_mut_slice(),
+                    fft_e.view_mut(),
+                    fft_scratch.as_mut_slice(),
+                    plans,
+                );
+                dest_fft.push(fft_e);
+                dest_fft.last().unwrap()
+            } else {
+                belief_from_var[*e].fft_to(
+                    fft_input_scratch.as_mut_slice(),
+                    fft_tmp.view_mut(),
+                    fft_scratch.as_mut_slice(),
+                    plans,
+                );
+                &fft_tmp
+            };
+            if acc_fft_init {
+                acc_fft *= v;
+            } else {
+                acc_fft.assign(&v);
+                acc_fft_init = true;
+            }
+            if clear_incoming {
+                belief_from_var[*e].reset();
+            }
+        }
+        // This could be done in O(l) instead of O(l^2) where l=dest.len() by
+        // better caching product computations.
+        let mut fft_scratch = plans.c2r.make_scratch_vec();
+        return (0..dest.len())
+            .map(move |i| {
+                let mut res = acc_fft.clone();
+                for (j, fft_op) in dest_fft.iter().enumerate() {
+                    if j != i {
+                        res *= fft_op;
+                    }
+                }
+                let mut acc = uniform_template.clone();
+                acc.ifft(res.view_mut(), fft_scratch.as_mut_slice(), plans);
+                acc.regularize();
+                acc
+            })
+            .collect::<Vec<_>>()
+            .into_iter();
+    }
+}
+
+fn factor_mul<'a>(
+    factor: &'a Factor,
+    belief_from_var: &'a mut EdgeSlice<Distribution>,
+    dest: &'a [VarId],
+    clear_incoming: bool,
+    pub_red: &'a PublicValue,
+) -> impl Iterator<Item = Distribution> + 'a {
+    // This is a simple, non-optimized algorithm.
+    let mut dest_iter = dest.iter();
+    std::iter::from_fn(move || {
+        if let Some(var) = dest_iter.next() {
+            // We compute the product and not a factor.
+            let is_product = factor.edges.get_index_of(var).unwrap() == 0;
+            let mut res: Option<Distribution> = None;
+            for (v, e) in factor.edges.iter() {
+                if v != var {
+                    res = Some(if let Some(res) = res {
+                        if is_product {
+                            res.op_multiply(&belief_from_var[*e])
+                        } else {
+                            res.op_multiply_factor(&belief_from_var[*e])
+                        }
+                    } else {
+                        belief_from_var[*e].take_or_clone(clear_incoming)
+                    });
+                }
+            }
+            Some(if is_product {
+                res.unwrap().op_multiply_cst(pub_red)
+            } else {
+                res.unwrap().op_multiply_cst_factor(pub_red)
+            })
+        } else {
+            if clear_incoming {
+                for e in factor.edges.values() {
+                    belief_from_var[*e].reset();
+                }
+            }
+            None
+        }
+    })
+}
+
+fn factor_lookup<'a>(
+    factor: &'a Factor,
+    belief_from_var: &'a mut EdgeSlice<Distribution>,
+    dest: &'a [VarId],
+    clear_incoming: bool,
+    table: &'a Table,
+) -> impl Iterator<Item = Distribution> + 'a {
+    // we know that there is no constant involved
+    assert_eq!(factor.edges.len(), 2);
+    dest.iter().map(move |dest| {
+        let i = factor.edges.get_index_of(dest).unwrap();
+        let distr = belief_from_var[factor.edges[1 - i]].clone();
+        let res = if i == 0 {
+            // dest is res
+            distr.map_table(table.values.as_slice())
+        } else {
+            distr.map_table_inv(table.values.as_slice())
+        };
+        if clear_incoming {
+            belief_from_var[factor.edges[1 - i]].reset();
+        }
+        res
+    })
+}
diff --git a/src/scalib_ext/scalib/src/sasca/bp_compute.rs b/src/scalib_ext/scalib/src/sasca/bp_compute.rs
new file mode 100644
index 00000000..4b804a96
--- /dev/null
+++ b/src/scalib_ext/scalib/src/sasca/bp_compute.rs
@@ -0,0 +1,628 @@
+use super::belief_propagation::{BPError, FftPlans};
+use super::factor_graph::PublicValue;
+use super::ClassVal;
+use ndarray::{azip, s, ArrayViewMut2, Zip};
+use realfft::num_complex::Complex;
+
+type Proba = f64;
+
+/// The minimum non-zero probability (to avoid denormalization, etc.)
+const MIN_PROBA: Proba = 1e-40;
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+enum DistrRepr {
+    Uniform,
+    Full(ndarray::Array2<Proba>),
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct Distribution {
+    multi: bool,
+    /// (nmulti, nc), same as DistrRepr::Full(x) => x.dim()
+    shape: (usize, usize),
+    value: DistrRepr,
+}
+
+impl Distribution {
+    pub fn multi(&self) -> bool {
+        self.multi
+    }
+    pub fn shape(&self) -> (usize, usize) {
+        self.shape
+    }
+    pub fn value(&self) -> Option<ndarray::ArrayView2<Proba>> {
+        if let DistrRepr::Full(v) = &self.value {
+            Some(v.view())
+        } else {
+            None
+        }
+    }
+    pub fn value_mut(&mut self) -> Option<ndarray::ArrayViewMut2<Proba>> {
+        if let DistrRepr::Full(v) = &mut self.value {
+            Some(v.view_mut())
+        } else {
+            None
+        }
+    }
+    pub fn from_array_single(array: ndarray::Array1<Proba>) -> Result<Self, BPError> {
+        let l = array.len();
+        if array.is_standard_layout() {
+            let array = array.into_shape((1, l)).expect("Non-contiguous array");
+            Ok(Self {
+                multi: false,
+                shape: array.dim(),
+                value: DistrRepr::Full(array),
+            })
+        } else {
+            Err(BPError::DistributionLayout(
+                array.shape().to_owned(),
+                array.strides().to_owned(),
+            ))
+        }
+    }
+    pub fn from_array_multi(array: ndarray::Array2<Proba>) -> Result<Self, BPError> {
+        if array.is_standard_layout() {
+            Ok(Self {
+                multi: true,
+                shape: array.dim(),
+                value: DistrRepr::Full(array),
+            })
+        } else {
+            Err(BPError::DistributionLayout(
+                array.shape().to_owned(),
+                array.strides().to_owned(),
+            ))
+        }
+    }
+    pub fn new_single(nc: usize) -> Self {
+        Self {
+            multi: false,
+            shape: (1, nc),
+            value: DistrRepr::Uniform,
+        }
+    }
+    pub fn new_multi(nc: usize, nmulti: u32) -> Self {
+        Self {
+            multi: true,
+            shape: (nmulti as usize, nc),
+            value: DistrRepr::Uniform,
+        }
+    }
+    pub fn new(multi: bool, nc: usize, nmulti: u32) -> Self {
+        if multi {
+            Self::new_multi(nc, nmulti)
+        } else {
+            Self::new_single(nc)
+        }
+    }
+    pub fn as_uniform(&self) -> Self {
+        Self {
+            multi: self.multi,
+            shape: self.shape,
+            value: DistrRepr::Uniform,
+        }
+    }
+    pub fn new_constant(&self, cst: &PublicValue) -> Self {
+        if let PublicValue::Multi(cst) = cst {
+            assert!(self.multi);
+            assert_eq!(self.shape.0, cst.len());
+        }
+        let mut value = ndarray::Array2::zeros(self.shape);
+        for (i, mut d) in value.axis_iter_mut(ndarray::Axis(0)).enumerate() {
+            let c = match cst {
+                PublicValue::Single(x) => *x,
+                PublicValue::Multi(x) => x[i],
+            };
+            d[c as usize] = 1.0;
+        }
+        Self {
+            multi: self.multi,
+            shape: self.shape,
+            value: DistrRepr::Full(value),
+        }
+    }
+    pub fn reset(&mut self) -> Self {
+        let mut new = self.as_uniform();
+        std::mem::swap(self, &mut new);
+        new
+    }
+
+    pub fn take_or_clone(&mut self, reset: bool) -> Self {
+        if reset {
+            self.reset()
+        } else {
+            self.clone()
+        }
+    }
+
+    pub fn multiply<'a>(&mut self, factors: impl Iterator<Item = &'a Distribution>) {
+        self.multiply_inner(factors, 0.0, 0.0);
+    }
+
+    pub fn multiply_reg<'a>(&mut self, factors: impl Iterator<Item = &'a Distribution>) {
+        self.multiply_inner(factors, MIN_PROBA, MIN_PROBA * MIN_PROBA);
+    }
+
+    fn multiply_inner<'a>(
+        &mut self,
+        factors: impl Iterator<Item = &'a Distribution>,
+        offset: Proba,
+        offset_prod: Proba,
+    ) {
+        for factor in factors {
+            assert_eq!(self.shape.1, factor.shape.1);
+            if self.multi & factor.multi {
+                assert_eq!(self.shape.0, factor.shape.0);
+            }
+            if let DistrRepr::Full(d) = &factor.value {
+                assert_eq!(d.dim(), factor.shape);
+                match (self.multi, factor.multi, &mut self.value) {
+                    (_, false, DistrRepr::Uniform) => {
+                        self.value = DistrRepr::Full(ndarray::Array::from_shape_fn(
+                            self.shape,
+                            |(_i, j)| d[(0, j)] + offset,
+                        ));
+                    }
+                    (false, true, DistrRepr::Uniform) => {
+                        let mut v = ndarray::Array2::ones(self.shape);
+                        for d in d.axis_iter(ndarray::Axis(0)) {
+                            azip!(v.slice_mut(s![0, ..]), d).for_each(|v, d| {
+                                *v = *v * (*d + offset) + offset_prod;
+                            });
+                        }
+                        self.value = DistrRepr::Full(v);
+                    }
+                    (true, true, DistrRepr::Uniform) => {
+                        self.value = DistrRepr::Full(d.map(|d| *d + offset));
+                    }
+                    (true, _, DistrRepr::Full(ref mut v)) => {
+                        azip!(v, d).for_each(|v, d| {
+                            *v = *v * (*d + offset) + offset_prod;
+                        });
+                    }
+                    (false, _, DistrRepr::Full(ref mut v)) => {
+                        for d in d.axis_iter(ndarray::Axis(0)) {
+                            azip!(v.slice_mut(s![0, ..]), d).for_each(|v, d| {
+                                *v = *v * (*d + offset) + offset_prod;
+                            });
+                        }
+                    }
+                }
+            }
+        }
+    }
+    pub fn divide_reg(state: &Distribution, div: &Distribution) -> Self {
+        let mut res = Self {
+            multi: state.multi | div.multi,
+            shape: (std::cmp::max(state.shape.0, div.shape.0), state.shape.1),
+            value: DistrRepr::Uniform,
+        };
+        assert!(res.shape == state.shape || state.shape == (1, res.shape.1));
+        assert!(res.shape == div.shape || div.shape == (1, res.shape.1));
+        let one = ndarray::Array2::ones((1, 1));
+        let (vst, vdiv) = match (&state.value, &div.value) {
+            (DistrRepr::Uniform, DistrRepr::Uniform) => {
+                return res;
+            }
+            (DistrRepr::Uniform, DistrRepr::Full(v)) => (&one, v),
+            (DistrRepr::Full(v), DistrRepr::Uniform) => (v, &one),
+            (DistrRepr::Full(vst), DistrRepr::Full(vdiv)) => (vst, vdiv),
+        };
+        res.value = DistrRepr::Full(
+            Zip::from(
+                vst.broadcast((std::cmp::max(vst.dim().0, vdiv.dim().0), vst.dim().1))
+                    .unwrap(),
+            )
+            .and_broadcast(vdiv)
+            .map_collect(|vst, vdiv| *vst / (*vdiv + MIN_PROBA)),
+        );
+        return res;
+    }
+    pub fn dividing_full(&mut self, other: &Distribution) {
+        match (&mut self.value, &other.value) {
+            (DistrRepr::Full(div), DistrRepr::Full(st)) => {
+                ndarray::azip!(div, st).for_each(|div, st| *div = *st / *div);
+            }
+            _ => {
+                unimplemented!();
+            }
+        }
+    }
+    pub fn not(&mut self) {
+        let inv_cst = (self.shape.1 - 1) as u32;
+        self.for_each_ignore(|mut d, _| {
+            xor_cst_slice(d.as_slice_mut().unwrap(), inv_cst);
+        });
+    }
+    pub fn wht(&mut self) {
+        self.for_each_error(|mut d, _| {
+            slice_wht(d.as_slice_mut().unwrap());
+        });
+    }
+    pub fn cumt(&mut self) {
+        self.for_each_error(|mut d, _| {
+            slice_cumt(d.as_slice_mut().unwrap());
+        });
+    }
+    pub fn cumti(&mut self) {
+        self.for_each_error(|mut d, _| {
+            slice_cumti(d.as_slice_mut().unwrap());
+        });
+    }
+    pub fn opandt(&mut self) {
+        self.for_each_error(|mut d, _| {
+            slice_opandt(d.as_slice_mut().unwrap());
+        });
+    }
+    pub(super) fn fft_to(
+        &self,
+        input_scratch: &mut [f64],
+        mut dest: ArrayViewMut2<Complex<f64>>,
+        fft_scratch: &mut [Complex<f64>],
+        plans: &FftPlans,
+    ) {
+        if let DistrRepr::Full(v) = &self.value {
+            for (distr, mut dest) in v.outer_iter().zip(dest.outer_iter_mut()) {
+                input_scratch.copy_from_slice(distr.as_slice().unwrap());
+                plans
+                    .r2c
+                    .process_with_scratch(input_scratch, dest.as_slice_mut().unwrap(), fft_scratch)
+                    .unwrap();
+            }
+        }
+    }
+    pub(super) fn ifft(
+        &mut self,
+        mut input: ArrayViewMut2<Complex<f64>>,
+        fft_scratch: &mut [Complex<f64>],
+        plans: &FftPlans,
+    ) {
+        self.ensure_full();
+        let mut v = self.value_mut().unwrap();
+        for (mut dest, mut input) in v.outer_iter_mut().zip(input.outer_iter_mut()) {
+            plans
+                .c2r
+                .process_with_scratch(
+                    input.as_slice_mut().unwrap(),
+                    dest.as_slice_mut().unwrap(),
+                    fft_scratch,
+                )
+                .unwrap();
+        }
+    }
+    pub fn is_full(&self) -> bool {
+        match &self.value {
+            DistrRepr::Full(_) => true,
+            DistrRepr::Uniform => false,
+        }
+    }
+    pub fn ensure_full(&mut self) {
+        if let DistrRepr::Uniform = self.value {
+            self.value = DistrRepr::Full(ndarray::Array2::from_elem(
+                self.shape,
+                1.0 / (self.shape.1 as Proba),
+            ));
+        }
+    }
+    pub fn full_zeros(&self) -> Self {
+        Self {
+            multi: self.multi,
+            shape: self.shape,
+            value: DistrRepr::Full(ndarray::Array2::zeros(self.shape)),
+        }
+    }
+    pub fn map_table(&self, table: &[ClassVal]) -> Self {
+        let mut value = ndarray::Array2::zeros(self.shape);
+        if let DistrRepr::Full(v_orig) = &self.value {
+            for (mut dest, orig) in value.outer_iter_mut().zip(v_orig.outer_iter()) {
+                for (i, p) in orig.iter().enumerate() {
+                    dest[table[i] as usize] += p;
+                }
+            }
+        } else {
+            // TODO optimize for bijective tables
+            for mut dest in value.outer_iter_mut() {
+                for i in 0..self.shape.1 {
+                    dest[table[i] as usize] += 1.0 / (self.shape.1 as f64);
+                }
+            }
+        }
+        Self {
+            shape: self.shape,
+            multi: self.multi,
+            value: DistrRepr::Full(value),
+        }
+    }
+    pub fn map_table_inv(&self, table: &[ClassVal]) -> Self {
+        let mut value = ndarray::Array2::zeros(self.shape);
+        if let DistrRepr::Full(v_orig) = &self.value {
+            for (mut dest, orig) in value.outer_iter_mut().zip(v_orig.outer_iter()) {
+                for (i, d) in dest.iter_mut().enumerate() {
+                    *d = orig[table[i] as usize];
+                }
+            }
+            Self {
+                shape: self.shape,
+                multi: self.multi,
+                value: DistrRepr::Full(value),
+            }
+        } else {
+            // here we are always uniform
+            Self {
+                shape: self.shape,
+                multi: self.multi,
+                value: DistrRepr::Uniform,
+            }
+        }
+    }
+    /// Normalize sum to one, and make values not too small
+    pub fn normalize(&mut self) {
+        self.for_each_ignore(|mut d, _| {
+            let norm_f = 1.0 / d.sum();
+            d.mapv_inplace(|x| x * norm_f);
+        })
+    }
+    /// 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);
+        })
+    }
+    pub fn make_non_zero_signed(&mut self) {
+        self.for_each_ignore(|mut d, _| {
+            d.mapv_inplace(|y| {
+                if y.is_sign_positive() {
+                    y.max(MIN_PROBA)
+                } else {
+                    y.min(-MIN_PROBA)
+                }
+            });
+        });
+    }
+    pub fn xor_cst(&mut self, cst: &PublicValue) {
+        self.for_each_ignore(|mut d, i| {
+            xor_cst_slice(d.as_slice_mut().unwrap(), cst.get(i));
+        });
+    }
+    pub fn and_cst(&mut self, cst: &PublicValue) {
+        self.for_each_ignore(|mut d, i| {
+            and_cst_slice(d.as_slice_mut().unwrap(), cst.get(i));
+        });
+    }
+    pub fn inv_and_cst(&mut self, cst: &PublicValue) {
+        self.for_each_ignore(|mut d, i| {
+            inv_and_cst_slice(d.as_slice_mut().unwrap(), cst.get(i));
+        });
+    }
+    pub fn add_cst(&mut self, cst: &PublicValue, sub: bool) {
+        let nc = self.shape.1;
+        if let DistrRepr::Full(v) = &mut self.value {
+            let mut tmp = vec![0.0f64; nc];
+            for (mut d, cst) in v.outer_iter_mut().zip(cst.iter(self.shape.0)) {
+                let d = d.as_slice_mut().unwrap();
+                let mut cst = usize::try_from(cst).unwrap() % nc;
+                if sub {
+                    cst = (nc - cst) % nc;
+                }
+                let size_first_block = self.shape.1 - cst;
+                tmp[cst..].copy_from_slice(&d[..size_first_block]);
+                tmp[..cst].copy_from_slice(&d[size_first_block..]);
+                d.copy_from_slice(tmp.as_slice());
+            }
+        }
+    }
+    pub fn op_multiply(&self, other: &Self) -> Self {
+        let mut res = self.full_zeros();
+        let nc = self.shape().1;
+        let u = 1.0f64 / (nc as f64);
+        for (k, mut res) in res.value_mut().unwrap().outer_iter_mut().enumerate() {
+            for i1 in 0..nc {
+                for i2 in 0..nc {
+                    let o = (((i1 * i2) as u32) % (nc as u32)) as usize;
+                    res[o] += self.value().map(|d| d[(k, i1)]).unwrap_or(u)
+                        * other.value().map(|d| d[(k, i2)]).unwrap_or(u);
+                }
+            }
+        }
+        res
+    }
+    pub fn op_multiply_factor(&self, other: &Self) -> Self {
+        let mut res = self.full_zeros();
+        let nc = self.shape().1;
+        let u = 1.0f64 / (nc as f64);
+        for (k, mut res) in res.value_mut().unwrap().outer_iter_mut().enumerate() {
+            for i1 in 0..nc {
+                for i2 in 0..nc {
+                    let o = (((i1 * i2) as u32) % (nc as u32)) as usize;
+                    res[i1] += self.value().map(|d| d[(k, o)]).unwrap_or(u)
+                        * other.value().map(|d| d[(k, i2)]).unwrap_or(u);
+                }
+            }
+        }
+        res
+    }
+    pub fn op_multiply_cst(&self, other: &PublicValue) -> Self {
+        let mut res = self.full_zeros();
+        let nc = self.shape().1;
+        let u = 1.0f64 / (nc as f64);
+        for (k, (mut res, i2)) in res
+            .value_mut()
+            .unwrap()
+            .outer_iter_mut()
+            .zip(other.iter(self.shape().0))
+            .enumerate()
+        {
+            let i2 = i2 as usize;
+            for i1 in 0..nc {
+                let o = (((i1 * i2) as u32) % (nc as u32)) as usize;
+                res[o] += self.value().map(|d| d[(k, i1)]).unwrap_or(u);
+            }
+        }
+        res
+    }
+    pub fn op_multiply_cst_factor(&self, other: &PublicValue) -> Self {
+        let mut res = self.full_zeros();
+        let nc = self.shape().1;
+        let u = 1.0f64 / (nc as f64);
+        for (k, (mut res, i2)) in res
+            .value_mut()
+            .unwrap()
+            .outer_iter_mut()
+            .zip(other.iter(self.shape().0))
+            .enumerate()
+        {
+            let i2 = i2 as usize;
+            for i1 in 0..nc {
+                let o = (((i1 * i2) as u32) % (nc as u32)) as usize;
+                res[i1] += self.value().map(|d| d[(k, o)]).unwrap_or(u);
+            }
+        }
+        res
+    }
+    pub fn for_each<F, G>(&mut self, mut f: F, default: G)
+    where
+        F: FnMut(ndarray::ArrayViewMut1<f64>, usize),
+        G: FnOnce(&mut Self),
+    {
+        if let DistrRepr::Full(v) = &mut self.value {
+            for (i, d) in v.axis_iter_mut(ndarray::Axis(0)).enumerate() {
+                f(d, i);
+            }
+        } else {
+            default(self);
+        }
+    }
+    pub fn for_each_ignore<F>(&mut self, f: F)
+    where
+        F: FnMut(ndarray::ArrayViewMut1<f64>, usize),
+    {
+        self.for_each(f, |_| {});
+    }
+    pub fn for_each_error<F>(&mut self, f: F)
+    where
+        F: FnMut(ndarray::ArrayViewMut1<f64>, usize),
+    {
+        self.for_each(f, |_| {
+            unimplemented!("This function must be called on Full distributions.");
+        });
+    }
+}
+
+/// Walsh-Hadamard transform (non-normalized).
+fn slice_wht(a: &mut [f64]) {
+    // The speed of this can be much improved, with the following techiques
+    // * improved memory locality
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let len = a.len();
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x + y;
+                a[j + h] = x - y;
+            }
+        }
+        h *= 2;
+    }
+}
+
+/// Cumulative transform (U in RLDA paper)
+fn slice_cumt(a: &mut [f64]) {
+    // Note: the speed of this can probably be much improved, with the following techiques
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let len = a.len();
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x + y;
+                a[j + h] = y;
+            }
+        }
+        h *= 2;
+    }
+}
+
+/// Cumulative inverse transform (U^-1 in RLDA paper)
+fn slice_cumti(a: &mut [f64]) {
+    // Note: the speed of this can probably be much improved, with the following techiques
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let len = a.len();
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x - y;
+                a[j + h] = y;
+            }
+        }
+        h *= 2;
+    }
+}
+
+/// Tansform for operand of AND (V in RLDA paper), involutive
+fn slice_opandt(a: &mut [f64]) {
+    // Note: the speed of this can probably be much improved, with the following techiques
+    // * use (auto-)vectorization
+    // * generate small static kernels
+    let len = a.len();
+    let mut h = 1;
+    while h < len {
+        for mut i in 0..(len / (2 * h) as usize) {
+            i *= 2 * h;
+            for j in i..(i + h) {
+                let x = a[j];
+                let y = a[j + h];
+                a[j] = x;
+                a[j + h] = x - y;
+            }
+        }
+        h *= 2;
+    }
+}
+
+fn xor_cst_slice(a: &mut [f64], cst: ClassVal) {
+    let leading_zeros = cst.leading_zeros();
+    if leading_zeros == 32 {
+        return;
+    }
+    let pivot_bit = ClassVal::BITS - 1 - leading_zeros;
+    let step = [1, 1 << (pivot_bit + 1) as usize];
+    let n_max = [1 << pivot_bit as usize, a.len()];
+    let n_iter = [n_max[0] / step[0], n_max[1] / step[1]];
+    let outer_i = if n_iter[0] < n_iter[1] { 0 } else { 1 };
+    for i in (0..n_max[outer_i]).step_by(step[outer_i]) {
+        for j in (0..n_max[1 - outer_i]).step_by(step[1 - outer_i]) {
+            let idx = i + j;
+            a.swap(idx, idx ^ cst as usize);
+        }
+    }
+}
+fn and_cst_slice(a: &mut [f64], cst: ClassVal) {
+    for i in 0..a.len() {
+        let j = i & (cst as usize);
+        if j != i {
+            a[j] += a[i];
+            a[i] = 0.0;
+        }
+    }
+}
+fn inv_and_cst_slice(a: &mut [f64], cst: ClassVal) {
+    for i in 0..a.len() {
+        let j = i & (cst as usize);
+        a[i] = a[j];
+    }
+}
diff --git a/src/scalib_ext/scalib/src/sasca/factor_graph.rs b/src/scalib_ext/scalib/src/sasca/factor_graph.rs
new file mode 100644
index 00000000..cc0644a2
--- /dev/null
+++ b/src/scalib_ext/scalib/src/sasca/factor_graph.rs
@@ -0,0 +1,397 @@
+use std::fmt;
+
+use indexmap::IndexMap;
+use thiserror::Error;
+
+use super::{ClassVal, NamedList};
+
+macro_rules! new_id {
+    ($it:ident, $vt:ident) => {
+        index_vec::define_index_type! {
+            pub struct $it = u32;
+            DISPLAY_FORMAT = "{}";
+        }
+        pub type $vt<T> = index_vec::IndexVec<$it, T>;
+        impl $it {
+            pub fn idx(self) -> usize {
+                <Self as Into<usize>>::into(self)
+            }
+            pub fn from_idx(x: usize) -> Self {
+                <Self as From<usize>>::from(x)
+            }
+        }
+    };
+}
+new_id!(VarId, VarVec);
+new_id!(FactorId, FactorVec);
+new_id!(EdgeId, EdgeVec);
+pub type EdgeSlice<T> = index_vec::IndexSlice<EdgeId, [T]>;
+pub type PublicId = usize;
+pub type TableId = usize;
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) struct Var {
+    pub(super) multi: bool,
+    pub(super) edges: IndexMap<FactorId, EdgeId>,
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) struct Factor {
+    pub(super) kind: FactorKind,
+    pub(super) multi: bool,
+    // res is first element (if there is a res), operands come next
+    pub(super) edges: IndexMap<VarId, EdgeId>,
+    // Is the result a variable (and not a public) ?
+    pub(super) has_res: bool,
+    // May not be allowed for all factor kinds
+    pub(super) publics: Vec<(PublicId, bool)>,
+}
+
+impl Factor {
+    pub(super) fn var_id(&self, var_order: usize) -> VarId {
+        *self.edges.get_index(var_order).unwrap().0
+    }
+    pub(super) fn res_id(&self) -> Option<VarId> {
+        self.has_res.then_some(self.var_id(0))
+    }
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) enum FactorKind<T = TableId> {
+    AND { vars_neg: Vec<bool> },
+    XOR,
+    NOT,
+    ADD,
+    MUL,
+    LOOKUP { table: T },
+}
+
+impl FactorKind {
+    fn merge(&self, a: ClassVal, b: ClassVal, nc: usize) -> ClassVal {
+        match self {
+            FactorKind::AND { vars_neg: _ } => a & b,
+            FactorKind::XOR => a ^ b,
+            FactorKind::ADD => (((a as u64) + (b as u64)) % (nc as u64)) as ClassVal,
+            FactorKind::MUL => (((a as u64) + (b as u64)) % (nc as u64)) as ClassVal,
+            FactorKind::NOT | FactorKind::LOOKUP { .. } => unreachable!(),
+        }
+    }
+    fn neutral(&self, nc: usize) -> ClassVal {
+        match self {
+            FactorKind::AND { vars_neg: _ } => (nc - 1) as ClassVal,
+            FactorKind::XOR | FactorKind::ADD => 0,
+            FactorKind::MUL => 1,
+            FactorKind::NOT | FactorKind::LOOKUP { .. } => unreachable!(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) struct Edge {
+    pub(super) var: VarId,
+    pub(super) pos_var: usize,
+    pub(super) factor: FactorId,
+    pub(super) pos_factor: usize,
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) struct Public {
+    pub(super) multi: bool,
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) struct Table {
+    pub(super) values: Vec<ClassVal>,
+}
+
+#[derive(Debug, Copy, Clone, serde::Serialize, serde::Deserialize)]
+pub(super) enum Node {
+    Var(VarId),
+    Factor(FactorId),
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct FactorGraph {
+    pub(super) nc: usize,
+    pub(super) vars: NamedList<Var>,
+    pub(super) factors: NamedList<Factor>,
+    pub(super) edges: EdgeVec<Edge>,
+    pub(super) publics: NamedList<Public>,
+    pub(super) tables: NamedList<Table>,
+    pub(super) petgraph: petgraph::Graph<Node, EdgeId, petgraph::Undirected>,
+    pub(super) var_graph_ids: VarVec<petgraph::graph::NodeIndex>,
+    pub(super) factor_graph_ids: FactorVec<petgraph::graph::NodeIndex>,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
+pub enum PublicValue {
+    Single(ClassVal),
+    Multi(Vec<ClassVal>),
+}
+impl PublicValue {
+    pub fn as_slice(&self) -> &[ClassVal] {
+        match self {
+            PublicValue::Single(x) => std::slice::from_ref(x),
+            PublicValue::Multi(x) => x.as_slice(),
+        }
+    }
+    pub fn iter(&self, n: usize) -> impl Iterator<Item = ClassVal> + '_ {
+        if let PublicValue::Multi(x) = self {
+            assert_eq!(x.len(), n);
+        }
+        (0..n).map(|i| self.get(i))
+    }
+    pub fn get(&self, n: usize) -> ClassVal {
+        match self {
+            PublicValue::Single(x) => *x,
+            PublicValue::Multi(x) => x[n],
+        }
+    }
+    pub fn map(&self, f: impl Fn(ClassVal) -> ClassVal) -> Self {
+        match self {
+            PublicValue::Single(x) => PublicValue::Single(f(*x)),
+            PublicValue::Multi(x) => PublicValue::Multi(x.iter().cloned().map(f).collect()),
+        }
+    }
+}
+impl fmt::Display for PublicValue {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            PublicValue::Single(x) => write!(f, "{}", x),
+            PublicValue::Multi(x) => write!(f, "{:?}", x.as_slice()),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Error)]
+pub enum FGError {
+    #[error("No variable named {0}.")]
+    NoVar(String),
+    #[error("No factor named {0}.")]
+    NoFactor(String),
+    #[error("No edge between variable {var} and factor {factor}.")]
+    NoEdge { var: String, factor: FactorId },
+    #[error("Failure at factor {0}. Expected result {1}, got {2}.")]
+    CheckFail(String, PublicValue, PublicValue),
+}
+
+type FGResult<T> = Result<T, FGError>;
+
+impl FactorGraph {
+    pub fn edge(&self, var: VarId, factor: FactorId) -> FGResult<EdgeId> {
+        self.var(var)
+            .edges
+            .get(&factor)
+            .map(|e| *e)
+            .ok_or_else(|| FGError::NoEdge {
+                var: self.vars.get_index(var.idx()).unwrap().0.to_owned(),
+                factor,
+            })
+    }
+    pub fn public_multi(&self) -> impl Iterator<Item = (&str, bool)> {
+        self.publics.iter().map(|(n, v)| (n.as_str(), v.multi))
+    }
+    pub fn get_varid(&self, var: &str) -> FGResult<VarId> {
+        self.vars
+            .get_index_of(var)
+            .map(VarId::from_idx)
+            .ok_or_else(|| FGError::NoVar(var.to_owned()))
+    }
+    pub(super) fn var(&self, var: VarId) -> &Var {
+        &self.vars[var.idx()]
+    }
+    pub fn var_multi(&self, var: VarId) -> bool {
+        self.var(var).multi
+    }
+    pub fn range_vars(&self) -> impl Iterator<Item = VarId> {
+        (0..self.vars.len()).map(VarId::from_idx)
+    }
+    pub fn range_factors(&self) -> impl Iterator<Item = FactorId> {
+        (0..self.factors.len()).map(FactorId::from_idx)
+    }
+    pub(super) fn factor(&self, factor: FactorId) -> &Factor {
+        &self.factors[factor.idx()]
+    }
+    pub fn get_factorid(&self, factor: &str) -> FGResult<FactorId> {
+        self.factors
+            .get_index_of(factor)
+            .map(FactorId::from_idx)
+            .ok_or_else(|| FGError::NoFactor(factor.to_owned()))
+    }
+    pub fn var_names(&self) -> impl Iterator<Item = &str> {
+        self.vars.keys().map(String::as_str)
+    }
+    pub fn vars(&self) -> impl Iterator<Item = (VarId, &str)> {
+        self.vars
+            .keys()
+            .enumerate()
+            .map(|(i, vn)| (VarId::from_idx(i), vn.as_str()))
+    }
+    pub fn var_name(&self, v: VarId) -> &str {
+        self.vars.get_index(v.idx()).unwrap().0.as_str()
+    }
+    pub fn factor_names(&self) -> impl Iterator<Item = &str> {
+        self.factors.keys().map(String::as_str)
+    }
+    pub fn factor_name(&self, f: FactorId) -> &str {
+        self.factors.get_index(f.idx()).unwrap().0.as_str()
+    }
+    pub fn factor_scope<'s>(&'s self, factor: FactorId) -> impl Iterator<Item = VarId> + 's {
+        self.factor(factor).edges.keys().cloned()
+    }
+    pub fn sanity_check(
+        &self,
+        public_values: Vec<PublicValue>,
+        var_assignments: VarVec<PublicValue>,
+    ) -> FGResult<()> {
+        assert_eq!(public_values.len(), self.publics.len());
+        assert_eq!(var_assignments.len(), self.vars.len());
+        let reduced_pub = self.reduce_pub(public_values.as_slice());
+        for ((factor_name, factor), cst) in self.factors.iter().zip(reduced_pub) {
+            let expected_res = factor
+                .res_id()
+                .map(|v_id| &var_assignments[v_id])
+                .unwrap_or(&cst);
+            let skip_res = if factor.has_res { 1 } else { 0 };
+            let mut ops = factor
+                .edges
+                .keys()
+                .skip(skip_res)
+                .map(|v_id| &var_assignments[*v_id]);
+            let res = match &factor.kind {
+                FactorKind::AND { vars_neg } => self.merge_pubs(
+                    &factor.kind,
+                    ops.zip(vars_neg.iter().cloned())
+                        .chain(std::iter::once((&cst, false))),
+                ),
+                FactorKind::XOR | FactorKind::ADD | FactorKind::MUL => self.merge_pubs(
+                    &factor.kind,
+                    ops.zip(std::iter::repeat(false))
+                        .chain(std::iter::once((&cst, false))),
+                ),
+                FactorKind::NOT => ops.next().unwrap().map(|x| self.not(x)),
+                FactorKind::LOOKUP { table } => ops
+                    .next()
+                    .unwrap()
+                    .map(|x| self.tables[*table].values[x as usize]),
+            };
+            if &res != expected_res {
+                return Err(FGError::CheckFail(
+                    factor_name.clone(),
+                    expected_res.clone(),
+                    res,
+                ));
+            }
+        }
+        Ok(())
+    }
+    pub(super) fn reduce_pub(&self, public_values: &[PublicValue]) -> FactorVec<PublicValue> {
+        self.factors
+            .values()
+            .map(|factor| {
+                match &factor.kind {
+                    // Not used
+                    FactorKind::NOT | FactorKind::LOOKUP { .. } => PublicValue::Single(0),
+                    _ => self.merge_pubs(
+                        &factor.kind,
+                        factor
+                            .publics
+                            .iter()
+                            .map(|(pub_id, nv)| (&public_values[*pub_id], *nv)),
+                    ),
+                }
+            })
+            .collect()
+    }
+    fn not(&self, x: ClassVal) -> ClassVal {
+        ((self.nc - 1) as ClassVal) ^ x
+    }
+    fn merge_pubs<'a>(
+        &self,
+        factor_kind: &FactorKind,
+        // bool is the "invserse" for every item
+        pubs: impl Iterator<Item = (&'a PublicValue, bool)>,
+    ) -> PublicValue {
+        let merge_inner = |p1: PublicValue, (p2, nv2): (&PublicValue, bool)| {
+            let f = |x: ClassVal, y: ClassVal| {
+                factor_kind.merge(x, if nv2 { self.not(y) } else { y }, self.nc)
+            };
+            match (p1, p2) {
+                (PublicValue::Single(c1), PublicValue::Single(c2)) => {
+                    PublicValue::Single(f(c1, *c2))
+                }
+                (PublicValue::Single(c1), PublicValue::Multi(c2)) => {
+                    PublicValue::Multi(c2.iter().map(|c2| f(c1, *c2)).collect())
+                }
+                (PublicValue::Multi(mut c1), PublicValue::Single(c2)) => {
+                    for c1 in c1.iter_mut() {
+                        *c1 = f(*c1, *c2);
+                    }
+                    PublicValue::Multi(c1)
+                }
+                (PublicValue::Multi(mut c1), PublicValue::Multi(c2)) => {
+                    for (c1, c2) in c1.iter_mut().zip(c2.iter()) {
+                        *c1 = f(*c1, *c2);
+                    }
+                    PublicValue::Multi(c1)
+                }
+            }
+        };
+        pubs.fold(
+            PublicValue::Single(factor_kind.neutral(self.nc)),
+            merge_inner,
+        )
+    }
+
+    pub(super) fn is_cyclic(&self, multi_exec: bool) -> bool {
+        if petgraph::algo::is_cyclic_undirected(&self.petgraph) {
+            return true;
+        }
+        if multi_exec {
+            return petgraph::algo::kosaraju_scc(&self.petgraph)
+                .into_iter()
+                .any(|scc| {
+                    scc.into_iter()
+                        .filter(|n| match self.petgraph[*n] {
+                            Node::Var(var_id) => {
+                                !self.vars.get_index(var_id.index()).unwrap().1.multi
+                            }
+                            Node::Factor(_) => false,
+                        })
+                        .count()
+                        > 1
+                });
+        } else {
+            return false;
+        }
+    }
+
+    pub(super) fn propagation_order(&self, var: VarId) -> Vec<(Node, Option<Node>)> {
+        let mut propagations = vec![(Node::Var(var), None)];
+        petgraph::visit::depth_first_search(&self.petgraph, [self.var_graph_ids[var]], |event| {
+            if let petgraph::visit::DfsEvent::TreeEdge(parent, node) = event {
+                propagations.push((self.petgraph[node], Some(self.petgraph[parent])));
+            }
+            petgraph::visit::Control::<()>::Continue
+        });
+        propagations.reverse();
+        propagations
+    }
+}
+
+impl Node {
+    pub(super) fn var(self) -> Option<VarId> {
+        if let Node::Var(id) = self {
+            Some(id)
+        } else {
+            None
+        }
+    }
+    pub(super) fn factor(self) -> Option<FactorId> {
+        if let Node::Factor(id) = self {
+            Some(id)
+        } else {
+            None
+        }
+    }
+}
diff --git a/src/scalib_ext/scalib/src/sasca/fg_build.rs b/src/scalib_ext/scalib/src/sasca/fg_build.rs
new file mode 100644
index 00000000..fda0247d
--- /dev/null
+++ b/src/scalib_ext/scalib/src/sasca/fg_build.rs
@@ -0,0 +1,299 @@
+use super::factor_graph as fg;
+use super::factor_graph::{FactorId, FactorVec, Node, VarVec};
+use super::fg_parser;
+use super::{ClassVal, NamedList, VarId};
+use indexmap::IndexMap;
+use thiserror::Error;
+
+#[derive(Debug, Clone, Error)]
+pub enum GraphBuildError {
+    #[error("Table {0} declared multiple times.")]
+    MultipleTableDecl(String),
+    #[error("Variable or public {0} declared multiple times.")]
+    MultipleVarDecl(String),
+    #[error("Multiple properties with name {0}, property name must be unique.")]
+    MultiplePropDecl(String),
+    #[error("Variable or public {0} not declared.")]
+    UnknownVar(String),
+    #[error("Table {0} not declared.")]
+    UnknownTable(String),
+    #[error("Operand {1} appears multiple times in factor {0}.")]
+    RepeatedOperand(String, String),
+    #[error("Constants appears as both result and operand in factor {0}.")]
+    CstOpRes(String),
+    #[error("Wrong length {1} for table {0}.")]
+    TableSize(String, usize),
+    #[error("Wrong value {1} for table {0}.")]
+    TableValue(String, ClassVal),
+    #[error("Value of table {0} not given.")]
+    MissingTableDef(String),
+    #[error("NC given more than once.")]
+    MultipleNc,
+    #[error("NC not given.")]
+    NoNc,
+    #[error("Could not parse graph description.\n{0}")]
+    Parse(String),
+}
+
+impl fg::FactorGraph {
+    fn build(nc: usize) -> Self {
+        Self {
+            nc,
+            vars: NamedList::new(),
+            factors: NamedList::new(),
+            edges: fg::EdgeVec::new(),
+            publics: NamedList::new(),
+            tables: NamedList::new(),
+            petgraph: petgraph::Graph::new_undirected(),
+            var_graph_ids: VarVec::new(),
+            factor_graph_ids: FactorVec::new(),
+        }
+    }
+    fn check_new_var(&self, name: &String) -> Result<(), GraphBuildError> {
+        if self.publics.contains_key(name) || self.vars.contains_key(name) {
+            return Err(GraphBuildError::MultipleVarDecl(name.clone()));
+        }
+        Ok(())
+    }
+    fn add_var(&mut self, name: String, multi: bool) -> Result<(), GraphBuildError> {
+        self.check_new_var(&name)?;
+        let (var_idx, _) = self.vars.insert_full(
+            name,
+            fg::Var {
+                multi,
+                edges: IndexMap::new(),
+            },
+        );
+        self.var_graph_ids
+            .push(self.petgraph.add_node(Node::Var(VarId::from_idx(var_idx))));
+        Ok(())
+    }
+    fn add_pub(&mut self, name: String, multi: bool) -> Result<(), GraphBuildError> {
+        self.check_new_var(&name)?;
+        self.publics.insert(name, fg::Public { multi });
+        Ok(())
+    }
+    fn add_table(&mut self, name: String, values: Vec<ClassVal>) -> Result<(), GraphBuildError> {
+        if self.tables.contains_key(&name) {
+            return Err(GraphBuildError::MultipleTableDecl(name));
+        }
+        if values.len() != self.nc {
+            return Err(GraphBuildError::TableSize(name, values.len()));
+        }
+        for v in values.iter() {
+            if (*v as usize) >= self.nc {
+                return Err(GraphBuildError::TableValue(name, *v));
+            }
+        }
+        self.tables.insert(name, fg::Table { values });
+        Ok(())
+    }
+    fn add_factor<'a>(
+        &mut self,
+        name: String,
+        kind: fg::FactorKind<&str>,
+        vars: impl Iterator<Item = &'a str>,
+    ) -> Result<(), GraphBuildError> {
+        if self.factors.contains_key(&name) {
+            return Err(GraphBuildError::MultiplePropDecl(name));
+        }
+        let factor_entry = self.factors.entry(name.clone());
+        let factor_id = FactorId::from_idx(factor_entry.index());
+        let mut edges = IndexMap::new();
+        let mut publics = Vec::new();
+        let mut multi = false;
+        let mut has_res = None;
+        let mut is_var = Vec::new();
+        for (i, var) in vars.enumerate() {
+            if let Some((var_id, _, v)) = self.vars.get_full_mut(var) {
+                has_res = Some(has_res.unwrap_or(true));
+                is_var.push(true);
+                let var_id = fg::VarId::from_idx(var_id);
+                let edge_id = fg::EdgeId::from_idx(self.edges.len());
+                if edges.insert(var_id, edge_id).is_some() {
+                    return Err(GraphBuildError::RepeatedOperand(name, var.to_owned()));
+                }
+                v.edges.insert(factor_id, edge_id);
+                self.edges.push(fg::Edge {
+                    var: var_id,
+                    pos_var: v.edges.len() - 1,
+                    factor: factor_id,
+                    pos_factor: edges.len() - 1,
+                });
+                multi |= v.multi;
+            } else if let Some((pub_id, _, public)) = self.publics.get_full(var) {
+                if has_res == Some(false) {
+                    return Err(GraphBuildError::CstOpRes(name));
+                }
+                has_res = Some(has_res.unwrap_or(false));
+                is_var.push(false);
+                publics.push((pub_id, kind.get_neg(i)));
+                multi |= public.multi;
+            } else {
+                return Err(GraphBuildError::UnknownVar(var.to_owned()));
+            }
+        }
+        let kind = kind
+            .map_vars_neg(|vn| {
+                vn.into_iter()
+                    .zip(is_var.into_iter())
+                    .filter(|(_, iv)| *iv)
+                    .map(|(vn, _)| vn)
+                    .collect()
+            })
+            .map_table(|t| {
+                self.tables
+                    .get_index_of(t)
+                    .ok_or_else(|| GraphBuildError::UnknownTable(t.to_owned()))
+            })?;
+        let factor = fg::Factor {
+            kind,
+            multi,
+            edges,
+            has_res: has_res.expect("at least one var"),
+            publics,
+        };
+        factor_entry.or_insert(factor);
+        self.factor_graph_ids
+            .push(self.petgraph.add_node(Node::Factor(factor_id)));
+        Ok(())
+    }
+    fn add_graph_edges(&mut self) {
+        for (i, e) in self.edges.iter_enumerated() {
+            self.petgraph.add_edge(
+                self.var_graph_ids[e.var],
+                self.factor_graph_ids[e.factor],
+                i,
+            );
+        }
+    }
+}
+impl<T> fg::FactorKind<T> {
+    fn map_table<S, E, F>(self, f: F) -> Result<fg::FactorKind<S>, E>
+    where
+        F: Fn(T) -> Result<S, E>,
+    {
+        Ok(match self {
+            fg::FactorKind::AND { vars_neg } => fg::FactorKind::AND { vars_neg },
+            fg::FactorKind::XOR => fg::FactorKind::XOR,
+            fg::FactorKind::NOT => fg::FactorKind::NOT,
+            fg::FactorKind::ADD => fg::FactorKind::ADD,
+            fg::FactorKind::MUL => fg::FactorKind::MUL,
+            fg::FactorKind::LOOKUP { table } => fg::FactorKind::LOOKUP { table: f(table)? },
+        })
+    }
+    fn map_vars_neg<F>(self, f: F) -> Self
+    where
+        F: FnOnce(Vec<bool>) -> Vec<bool>,
+    {
+        match self {
+            fg::FactorKind::AND { vars_neg } => fg::FactorKind::AND {
+                vars_neg: f(vars_neg),
+            },
+            x => x,
+        }
+    }
+    fn get_neg(&self, i: usize) -> bool {
+        match self {
+            fg::FactorKind::AND { vars_neg } => vars_neg[i],
+            _ => false,
+        }
+    }
+}
+impl fg_parser::Expr {
+    fn as_factor_kind(&self) -> fg::FactorKind<&str> {
+        fn get_neg(vars: &Vec<fg_parser::NVar>, neg: bool) -> Vec<bool> {
+            // Include the result of the operation
+            std::iter::once(neg)
+                .chain(vars.iter().map(|v| v.neg ^ neg))
+                .collect()
+        }
+        match self {
+            Self::Not(_) => fg::FactorKind::NOT,
+            Self::Lookup { table, .. } => fg::FactorKind::LOOKUP {
+                table: table.as_str(),
+            },
+            Self::Add(_) => fg::FactorKind::ADD,
+            Self::Mul(_) => fg::FactorKind::MUL,
+            Self::Xor(_) => fg::FactorKind::XOR,
+            Self::And(vars) => fg::FactorKind::AND {
+                vars_neg: get_neg(vars, false),
+            },
+            // Use De Morgan's law to convert OR to AND.
+            Self::Or(vars) => fg::FactorKind::AND {
+                vars_neg: get_neg(vars, true),
+            },
+        }
+    }
+    fn vars(&self) -> impl Iterator<Item = &str> {
+        match self {
+            Self::Not(var) | Self::Lookup { var, .. } => vec![var.0.as_str()],
+            Self::Xor(v) | Self::Add(v) | Self::Mul(v) => v.iter().map(|v| v.0.as_str()).collect(),
+            Self::And(v) | Self::Or(v) => v.iter().map(|v| v.var.0.as_str()).collect(),
+        }
+        .into_iter()
+    }
+}
+pub(super) fn build_graph(
+    stmts: &[fg_parser::Statement],
+    mut tables: std::collections::HashMap<String, Vec<ClassVal>>,
+) -> Result<fg::FactorGraph, GraphBuildError> {
+    let nc = get_nc(stmts)?;
+    let mut graph = fg::FactorGraph::build(nc as usize);
+    for s in stmts {
+        match s {
+            fg_parser::Statement::VarDecl(vd) => {
+                graph.add_var(vd.name.0.clone(), vd.multi)?;
+            }
+            fg_parser::Statement::PubDecl(vd) => {
+                graph.add_pub(vd.name.0.clone(), vd.multi)?;
+            }
+            fg_parser::Statement::TableDecl { name, val } => {
+                let val = if let Some(val) = val {
+                    val.clone()
+                } else if let Some(val) = tables.get_mut(name) {
+                    std::mem::replace(val, Vec::new())
+                } else {
+                    return Err(GraphBuildError::MissingTableDef(name.clone()));
+                };
+                graph.add_table(name.clone(), val)?;
+            }
+            fg_parser::Statement::Invalid
+            | fg_parser::Statement::Empty
+            | fg_parser::Statement::Property { .. }
+            | fg_parser::Statement::NC(_) => {}
+        }
+    }
+    let mut anon_names = (0..).map(|i| format!("ANONYMOUS_{}", i));
+    for s in stmts {
+        if let fg_parser::Statement::Property { name, dest, expr } = s {
+            let name = name.clone().unwrap_or_else(|| anon_names.next().unwrap());
+            graph.add_factor(
+                name,
+                expr.as_factor_kind(),
+                std::iter::once(dest.0.as_str()).chain(expr.vars()),
+            )?;
+        }
+    }
+    graph.add_graph_edges();
+    Ok(graph)
+}
+
+fn get_nc(stmts: &[fg_parser::Statement]) -> Result<u64, GraphBuildError> {
+    let mut nc_decls = stmts.iter().filter_map(|s| {
+        if let fg_parser::Statement::NC(nc) = s {
+            Some(*nc)
+        } else {
+            None
+        }
+    });
+    if let Some(nc) = nc_decls.next() {
+        if nc_decls.next().is_some() {
+            return Err(GraphBuildError::MultipleNc);
+        } else {
+            Ok(nc)
+        }
+    } else {
+        return Err(GraphBuildError::NoNc);
+    }
+}
diff --git a/src/scalib_ext/scalib/src/sasca/fg_parser.rs b/src/scalib_ext/scalib/src/sasca/fg_parser.rs
new file mode 100644
index 00000000..82358649
--- /dev/null
+++ b/src/scalib_ext/scalib/src/sasca/fg_parser.rs
@@ -0,0 +1,171 @@
+use super::ClassVal;
+use ariadne::{Color, Fmt, Label, Report, ReportKind, Source};
+use chumsky::prelude::*;
+
+#[derive(Debug, Clone)]
+pub(super) struct Var(pub(super) String);
+
+#[derive(Debug, Clone)]
+pub(super) struct NVar {
+    pub(super) var: Var,
+    pub(super) neg: bool,
+}
+
+impl NVar {
+    fn new(var: Var, neg: bool) -> Self {
+        Self { var, neg }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(super) enum Expr {
+    Not(Var),
+    Lookup { var: Var, table: String },
+    Add(Vec<Var>),
+    Mul(Vec<Var>),
+    Xor(Vec<Var>),
+    And(Vec<NVar>),
+    Or(Vec<NVar>),
+}
+
+#[derive(Debug, Clone)]
+pub(super) struct VarDecl {
+    pub(super) name: Var,
+    pub(super) multi: bool,
+}
+
+#[derive(Debug, Clone)]
+pub(super) enum Statement {
+    Invalid,
+    Empty,
+    Property {
+        name: Option<String>,
+        dest: Var,
+        expr: Expr,
+    },
+    NC(u64),
+    VarDecl(VarDecl),
+    PubDecl(VarDecl),
+    TableDecl {
+        name: String,
+        val: Option<Vec<ClassVal>>,
+    },
+}
+
+fn parser() -> impl Parser<char, Vec<Statement>, Error = Simple<char>> {
+    let pad = just(' ').or(just('\t')).repeated();
+    let space = pad.at_least(1);
+    let op = |c| just(c).delimited_by(pad, pad);
+    let ident = pad.ignore_then(text::ident()).then_ignore(pad);
+    let kw = |s| text::keyword::<_, _, Simple<char>>(s).delimited_by(pad, space);
+    let var = ident.map(|s| Var(s));
+    let not_var = || op('!').ignore_then(var);
+    let nvar = || {
+        var.map(|v| NVar::new(v, false))
+            .or(not_var().map(|v| NVar::new(v, true)))
+    };
+    let op_nexpr = |c, f| nvar().separated_by(op(c)).at_least(2).map(f);
+    let op_expr = |c, f| var.separated_by(op(c)).at_least(2).map(f);
+    let expr = ident
+        .then(var.delimited_by(op('['), op(']')))
+        .map(|(table, var)| Expr::Lookup { table, var })
+        .or(op_expr('^', Expr::Xor as fn(_) -> _))
+        .or(op_nexpr('&', Expr::And as fn(_) -> _))
+        .or(op_nexpr('|', Expr::Or as fn(_) -> _))
+        .or(op_expr('+', Expr::Add as fn(_) -> _))
+        .or(op_expr('*', Expr::Mul as fn(_) -> _))
+        .or(not_var().map(|v| Expr::Not(v)));
+    let prop = kw("PROPERTY")
+        .ignore_then(ident.then_ignore(op(':')).or_not())
+        .then(var)
+        .then_ignore(op('='))
+        .then(expr)
+        .map(|((name, dest), expr)| Statement::Property { name, dest, expr });
+    let nc = kw("NC")
+        .ignore_then(text::int(10))
+        .map(|nc: String| Statement::NC(nc.parse().unwrap()));
+    let comment = op('#').then_ignore(filter(|c| *c != '\n' && *c != '\r').repeated());
+    let var_decl = kw("VAR")
+        .ignore_then(kw("SINGLE").to(false).or(kw("MULTI").to(true)))
+        .then(var)
+        .map(|(multi, name)| Statement::VarDecl(VarDecl { name, multi }));
+    let pub_decl = kw("PUB")
+        .ignore_then(kw("SINGLE").to(false).or(kw("MULTI").to(true)))
+        .then(var)
+        .map(|(multi, name)| Statement::PubDecl(VarDecl { name, multi }));
+    let table_val = text::int(10)
+        .map(|x: String| x.parse().unwrap())
+        .separated_by(op(','))
+        .allow_trailing()
+        .delimited_by(op('['), op(']'));
+    let table = kw("TABLE")
+        .ignore_then(ident)
+        .then(op('=').ignore_then(table_val).or_not())
+        .map(|(name, val)| Statement::TableDecl { name, val });
+    let graph = prop
+        .or(nc)
+        .or(var_decl)
+        .or(pub_decl)
+        .or(table)
+        .or(pad.at_least(0).to(Statement::Empty))
+        .then_ignore(comment.or_not())
+        .recover_with(skip_until(['\n', '\r'], |_| Statement::Invalid))
+        .separated_by(text::newline())
+        .allow_leading()
+        .allow_trailing()
+        .then_ignore(end());
+    graph
+}
+
+/// Parse the factor graph description in src, and returns the statements if no
+/// error, otherwise the error is a locale-encoded error message.
+pub(super) fn parse(src: &str) -> Result<Vec<Statement>, Vec<u8>> {
+    let (graph, errs) = parser().parse_recovery_verbose(src);
+    let err = !errs.is_empty();
+    let mut err_str = Vec::new();
+    for e in errs {
+        let msg = if let chumsky::error::SimpleReason::Custom(msg) = e.reason() {
+            msg.clone()
+        } else {
+            format!(
+                "Unexpected token, expected {}",
+                if e.expected().len() == 0 {
+                    "something else".to_string()
+                } else {
+                    e.expected()
+                        .map(|expected| match expected {
+                            Some(expected) => expected.to_string(),
+                            None => "end of input".to_string(),
+                        })
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                },
+            )
+        };
+        let report = Report::build(ReportKind::Error, (), e.span().start)
+            .with_code(3)
+            .with_message(msg)
+            .with_label(
+                Label::new(e.span())
+                    .with_message(match e.reason() {
+                        chumsky::error::SimpleReason::Custom(msg) => msg.clone(),
+                        _ => format!(
+                            "Unexpected {}",
+                            e.found()
+                                .map(|c| format!("token {}", c.fg(Color::Red)))
+                                .unwrap_or_else(|| "end of input".to_string())
+                        ),
+                    })
+                    .with_color(Color::Red),
+            );
+        report
+            .finish()
+            .write(Source::from(&src), &mut err_str)
+            .unwrap();
+    }
+    if err {
+        return Err(err_str);
+    } else {
+        return Ok(graph.unwrap());
+    }
+}
diff --git a/src/scalib_ext/scalib/src/sasca/mod.rs b/src/scalib_ext/scalib/src/sasca/mod.rs
new file mode 100644
index 00000000..cb5d5d1b
--- /dev/null
+++ b/src/scalib_ext/scalib/src/sasca/mod.rs
@@ -0,0 +1,23 @@
+mod belief_propagation;
+mod bp_compute;
+mod factor_graph;
+mod fg_build;
+mod fg_parser;
+
+pub use bp_compute::Distribution;
+
+pub type ClassVal = u32;
+type NamedList<T> = indexmap::IndexMap<String, T>;
+
+pub use belief_propagation::{BPError, BPState};
+pub use factor_graph::{EdgeId, FactorGraph, FactorId, PublicValue, VarId};
+pub use fg_build::GraphBuildError;
+
+pub fn build_graph(
+    description: &str,
+    tables: std::collections::HashMap<String, Vec<ClassVal>>,
+) -> Result<FactorGraph, GraphBuildError> {
+    let stmts = fg_parser::parse(description)
+        .map_err(|e| GraphBuildError::Parse(String::from_utf8_lossy(e.as_slice()).into_owned()))?;
+    fg_build::build_graph(stmts.as_slice(), tables)
+}


Test patch:
diff --git a/src/scalib_ext/scalib/src/mttest.rs b/src/scalib_ext/scalib/src/mttest.rs
index 3e31e470..5306930d 100644
--- a/src/scalib_ext/scalib/src/mttest.rs
+++ b/src/scalib_ext/scalib/src/mttest.rs
@@ -52,7 +52,7 @@ impl MultivarCSAcc {
     /// It allows to estimate for a class c the central sums (CS) such that
     /// at index j such that for the set combi:
     ///
-    ///     cs_Q[c,index(combi),j] = sum_i(prod_{s \in combi} Q[c,pois[s,j]])
+    /// cs_Q[c,index(combi),j] = sum_i(prod_{s \in combi} Q[c,pois[s,j]])
     ///
     /// for which:
     ///    - combi: is a unique combination in the multiset (0..d)U(0..d). s are the elements in S.
diff --git a/src/scalib_ext/scalib/tests/lda.rs b/src/scalib_ext/scalib/tests/lda.rs
index f6babd29..74d2cf02 100644
--- a/src/scalib_ext/scalib/tests/lda.rs
+++ b/src/scalib_ext/scalib/tests/lda.rs
@@ -4,7 +4,6 @@ use ndarray_rand::RandomExt;
 use ndarray_stats::QuantileExt;
 use scalib::lda;
 
-extern crate openblas_src as _;
 #[test]
 fn fit_lda() {
     let ns = 100;
@@ -12,11 +11,9 @@ fn fit_lda() {
     let nc = 16;
     let p = 4;
 
-    let mut lda = lda::LDA::new(nc, p, ns);
     let x = Array2::<i16>::random((n, ns), Uniform::new(0, 100));
     let y = Array1::<u16>::random(n, Uniform::new(0, nc as u16));
-
-    lda.fit(x.view(), y.view(), 0);
+    let lda = lda::LdaAcc::new(nc, x.view(), y.view(), 0);
 }
 
 #[test]
@@ -26,11 +23,9 @@ fn predict_lda() {
     let nc = 16;
     let p = 4;
 
-    let mut lda = lda::LDA::new(nc, p, ns);
     let x = Array2::<i16>::random((n, ns), Uniform::new(0, 100));
     let y = Array1::<u16>::random(n, Uniform::new(0, nc as u16));
-
-    lda.fit(x.view(), y.view(), 0);
+    let lda = lda::LdaAcc::new(nc, x.view(), y.view(), 0).lda(p).unwrap();
 
     lda.predict_proba(x.view());
 }
@@ -42,7 +37,6 @@ fn predict_lda_functional() {
     let nc = 16;
     let p = 4;
 
-    let mut lda = lda::LDA::new(nc, p, ns);
     let noise = Array2::<i16>::random((n, ns), Uniform::new(-10, 10));
     let mut signal = Array2::<i16>::zeros((n, ns));
     let y = Array1::<u16>::random(n, Uniform::new(0, nc as u16));
@@ -52,7 +46,7 @@ fn predict_lda_functional() {
         .for_each(|(mut s, y)| s.fill((y * 100) as i16));
     let x = &noise + signal;
 
-    lda.fit(x.view(), y.view(), 0);
+    let lda = lda::LdaAcc::new(nc, x.view(), y.view(), 0).lda(p).unwrap();
     let prs = lda.predict_proba(x.view());
 
     prs.outer_iter().zip(y.iter()).for_each(|(prs, y)| {
diff --git a/tests/test_factorgraph.py b/tests/test_factorgraph.py
new file mode 100644
index 00000000..1d47e7ef
--- /dev/null
+++ b/tests/test_factorgraph.py
@@ -0,0 +1,698 @@
+import pytest
+from scalib.attacks import FactorGraph, BPState
+import numpy as np
+import os
+import copy
+
+
+def normalize_distr(x):
+    return x / x.sum(axis=-1, keepdims=True)
+
+
+def make_distri(nc, n):
+    return normalize_distr(np.random.randint(1, 10000000, (n, nc)).astype(np.float64))
+
+
+def test_table():
+    """
+    Test Table lookup
+    """
+    nc = 16
+    n = 100
+    table = np.random.permutation(nc).astype(np.uint32)
+    distri_x = np.random.randint(1, 2048, (n, nc))
+    distri_x = (distri_x.T / np.sum(distri_x, axis=1)).T
+
+    graph = f"""
+            PROPERTY y = table[x]
+            TABLE table
+            VAR MULTI x
+            VAR MULTI y
+            NC {nc}
+            """
+    graph = FactorGraph(graph, {"table": table})
+    x = distri_x.argmax(axis=1)
+    graph.sanity_check({}, {"x": x, "y": table[x]})
+    bp_state = BPState(graph, n)
+
+    bp_state.set_evidence("x", distri_x)
+
+    bp_state.bp_loopy(1, True)
+    distri_y = bp_state.get_distribution("y")
+
+    distri_y_ref = np.zeros(distri_x.shape)
+    for x in range(nc):
+        y = table[x]
+        distri_y_ref[:, y] = distri_x[:, x]
+
+    assert np.allclose(distri_y_ref, distri_y)
+
+
+def test_table_non_bij():
+    """
+    Test non-bijective Table lookup
+    """
+    nc = 2
+    n = 1
+    table = np.array([0, 0], dtype=np.uint32)
+    distri_x = np.array([[0.5, 0.5]])
+    distri_y = np.array([[0.8, 0.2]])
+
+    graph = f"""
+            PROPERTY y = table[x]
+            TABLE table
+            VAR MULTI x
+            VAR MULTI y
+            NC {nc}
+            """
+    graph = FactorGraph(graph, {"table": table})
+
+    graph.sanity_check({}, {"x": np.array([0, 1]), "y": np.array([0, 0])})
+
+    bp_state = BPState(graph, n)
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+
+    bp_state.bp_loopy(1, True)
+    distri_x = bp_state.get_distribution("x")
+    distri_y = bp_state.get_distribution("y")
+
+    distri_x_ref = np.array([0.5, 0.5])
+    distri_y_ref = np.array([1.0, 0.0])
+
+    assert np.allclose(distri_x_ref, distri_x)
+    assert np.allclose(distri_y_ref, distri_y)
+
+
+def test_not():
+    """
+    Test NOT operation
+    """
+    nc = 2
+    n = 1
+    distri_x = np.array([[0.6, 0.4]])
+    distri_y = np.array([[0.8, 0.2]])
+
+    graph = f"""
+            PROPERTY y = !x
+            VAR MULTI x
+            VAR MULTI y
+            NC {nc}
+            """
+    graph = FactorGraph(graph)
+
+    graph.sanity_check({}, {"x": np.array([0, 1]), "y": np.array([1, 0])})
+
+    bp_state = BPState(graph, n)
+
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+
+    bp_state.bp_loopy(1, True)
+    distri_x_bp = bp_state.get_distribution("x")
+    distri_y_bp = bp_state.get_distribution("y")
+
+    t = np.array([1, 0])
+    distri_x_ref = normalize_distr(distri_x * distri_y[0, t][np.newaxis, :])
+    distri_y_ref = normalize_distr(distri_y * distri_x[0, t][np.newaxis, :])
+
+    assert np.allclose(distri_x_ref, distri_x_bp)
+    assert np.allclose(distri_y_ref, distri_y_bp)
+
+
+def test_and_public():
+    """
+    Test AND with public data
+    """
+    nc = 16
+    n = 100
+    public = np.random.randint(0, nc, n, dtype=np.uint32)
+    distri_x = np.random.randint(1, 100, (n, nc))
+    distri_x = (distri_x.T / np.sum(distri_x, axis=1)).T
+
+    graph = f"""
+        # some comments
+        NC {nc}
+        PROPERTY y = x & p
+        VAR MULTI y
+        VAR MULTI x
+        PUB MULTI p#come comments
+        """
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n, {"p": public})
+    bp_state.set_evidence("x", distri_x)
+
+    bp_state.bp_loopy(2, False)
+
+    distri_y = bp_state.get_distribution("y")
+    distri_y_ref = np.zeros(distri_x.shape)
+    for x in range(nc):
+        y = x & public
+        distri_y_ref[np.arange(n), y] += distri_x[np.arange(n), x]
+
+    assert np.allclose(distri_y_ref, distri_y)
+
+    bp_state2 = BPState(graph, n, {"p": public})
+    bp_state2.set_evidence("x", distri_x)
+    bp_state2.bp_acyclic("y")
+    assert np.allclose(distri_y_ref, bp_state2.get_distribution("y"))
+
+
+def test_xor_public():
+    """
+    Test XOR with public data
+    """
+    nc = 16
+    n = 100
+    public = np.random.randint(0, nc, n, dtype=np.uint32)
+    public2 = np.random.randint(0, nc, n, dtype=np.uint32)
+    distri_x = np.random.randint(1, 100, (n, nc))
+    distri_x = (distri_x.T / np.sum(distri_x, axis=1)).T
+
+    graph = f"""
+        PROPERTY y = x ^ p ^ p2
+        VAR MULTI y
+        VAR MULTI x
+        PUB MULTI p
+        PUB MULTI p2
+        NC {nc}
+        """
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n, {"p": public, "p2": public2})
+    bp_state.set_evidence("x", distri_x)
+
+    bp_state.bp_loopy(2, False)
+
+    distri_y = bp_state.get_distribution("y")
+    distri_y_ref = np.zeros(distri_x.shape)
+    for x in range(nc):
+        y = x ^ public ^ public2
+        distri_y_ref[np.arange(n), y] = distri_x[np.arange(n), x]
+
+    distri_x = bp_state.get_distribution("x")
+
+    print("p ^ p2")
+    print(public ^ public2)
+    print("y_ref")
+    print(distri_y_ref)
+    print("y")
+    print(distri_y)
+    print("x")
+    print(distri_x)
+    assert distri_x is not None
+    assert distri_y is not None
+    assert np.allclose(distri_y_ref, distri_y)
+
+
+def test_AND():
+    """
+    Test AND between distributions
+    """
+    cases = [
+        (
+            np.array([[0.5, 0.5]]),  # uniform x
+            np.array([[0.0, 1.0]]),  # y == 1
+            np.array([[0.0, 1.0]]),  # z == 1
+        ),
+        (
+            np.array([[0.1, 0.9]]),  # x == 1
+            np.array([[0.5, 0.5]]),  # uniform y
+            np.array([[1.0, 0.0]]),  # z == 1
+        ),
+        (
+            np.array([[0.1, 0.9]]),  # x == 1
+            np.array([[0.4, 0.6]]),  # uniform y
+            np.array([[0.8, 0.2]]),  # z == 1
+        ),
+        (make_distri(2, 1) for _ in range(3)),
+        (make_distri(4, 1) for _ in range(3)),
+        (make_distri(256, 4) for _ in range(3)),
+    ]
+
+    for distri_x, distri_y, distri_z in cases:
+        print("#### Test case:")
+        print(distri_x)
+        print(distri_y)
+        print(distri_z)
+        n, nc = distri_x.shape
+        graph = f"""
+            # some comments
+            NC {nc}
+            PROPERTY z = x&y
+            VAR MULTI z
+            VAR MULTI x
+            VAR MULTI y
+
+            """
+        graph = FactorGraph(graph)
+        bp_state = BPState(graph, n)
+        bp_state.set_evidence("x", distri_x)
+        bp_state.set_evidence("y", distri_y)
+        bp_state.set_evidence("z", distri_z)
+
+        distri_x_ref = np.zeros(distri_x.shape)
+        distri_y_ref = np.zeros(distri_y.shape)
+        distri_z_ref = np.zeros(distri_z.shape)
+
+        for x in range(nc):
+            for y in range(nc):
+                distri_x_ref[:, x] += distri_z[:, x & y] * distri_y[:, y]
+                distri_y_ref[:, y] += distri_z[:, x & y] * distri_x[:, x]
+                distri_z_ref[:, x & y] += distri_x[:, x] * distri_y[:, y]
+
+        distri_x_ref = normalize_distr(distri_x_ref * distri_x)
+        distri_y_ref = normalize_distr(distri_y_ref * distri_y)
+        distri_z_ref = normalize_distr(distri_z_ref * distri_z)
+
+        print("#### Ref:")
+        print(distri_x_ref)
+        print(distri_y_ref)
+        print(distri_z_ref)
+
+        bp_state.bp_loopy(1, True)
+        distri_x = bp_state.get_distribution("x")
+        distri_y = bp_state.get_distribution("y")
+        distri_z = bp_state.get_distribution("z")
+
+        print("#### Got:")
+        print(distri_x)
+        print(distri_y)
+        print(distri_z)
+
+        assert np.allclose(distri_z_ref, distri_z)
+        assert np.allclose(distri_x_ref, distri_x)
+        assert np.allclose(distri_y_ref, distri_y)
+
+
+def test_and_not_or():
+    nc = 16
+    n = 100
+    nc = 2
+    n = 1
+    p = np.random.randint(0, nc, n, dtype=np.uint32)
+    t = np.random.randint(0, nc, n, dtype=np.uint32)
+    nt = (nc - 1) ^ t
+    distri_x = make_distri(nc, n)
+    distri_y = make_distri(nc, n)
+    distri_z = make_distri(nc, n)
+
+    graph = f"""
+        NC {nc}
+        VAR MULTI x
+        VAR MULTI y
+        VAR MULTI z
+        PUB MULTI p
+        PUB MULTI t
+        PROPERTY z = !y & x & p & !t
+        """
+    graph2 = f"""
+        NC {nc}
+        VAR MULTI x
+        VAR MULTI y
+        VAR MULTI ny
+        PROPERTY ny = !y
+        VAR MULTI z
+        VAR MULTI nz
+        PROPERTY z = !nz
+        PUB MULTI p
+        PUB MULTI nt
+        VAR MULTI a
+        VAR MULTI na
+        PROPERTY na = !a
+        VAR MULTI b
+        VAR MULTI nb
+        PROPERTY nb = !b
+        PROPERTY a = x & p
+        PROPERTY b = nt & ny 
+        PROPERTY nz = na | nb 
+        """
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n, {"p": p, "t": t})
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+    bp_state.set_evidence("z", distri_z)
+    bp_state.bp_loopy(1, True)
+    import sys
+
+    print("BP2", file=sys.stderr)
+
+    graph2 = FactorGraph(graph2)
+    bp_state2 = BPState(graph2, n, {"p": p, "nt": nt})
+    bp_state2.set_evidence("x", distri_x)
+    bp_state2.set_evidence("y", distri_y)
+    bp_state2.set_evidence("z", distri_z)
+    bp_state2.bp_loopy(8, False)
+
+    for v in ["x", "y", "z"]:
+        print(v)
+        d = bp_state.get_distribution(v)
+        d2 = bp_state2.get_distribution(v)
+        assert np.allclose(d, d2)
+
+    bp_state3 = BPState(graph, n, {"p": p, "t": t})
+    bp_state3.set_evidence("x", distri_x)
+    bp_state3.set_evidence("y", distri_y)
+    bp_state3.set_evidence("z", distri_z)
+    for v in ["x", "y", "z"]:
+        print(v)
+        bp_state3.bp_acyclic(v)
+        assert np.allclose(bp_state2.get_distribution(v), bp_state3.get_distribution(v))
+
+
+def test_ADD():
+    """
+    Test ADD between distributions
+    """
+    nc = 251
+    n = 4
+    distri_x = make_distri(nc, n)
+    distri_y = make_distri(nc, n)
+
+    graph = f"""
+        # some comments
+        NC {nc}
+        PROPERTY z = x+y
+        VAR MULTI z
+        VAR MULTI x
+        VAR MULTI y
+
+        """
+
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n)
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+
+    bp_state.bp_loopy(1, True)
+    distri_z = bp_state.get_distribution("z")
+
+    distri_z_ref = np.zeros(distri_z.shape)
+    msg = np.zeros(distri_z.shape)
+
+    for x in range(nc):
+        for y in range(nc):
+            distri_z_ref[:, (x + y) % nc] += distri_x[:, x] * distri_y[:, y]
+
+    distri_z_ref = (distri_z_ref.T / np.sum(distri_z_ref, axis=1)).T
+    assert np.allclose(distri_z_ref, distri_z)
+
+
+def test_add_cst():
+    """
+    Test ADD of a constant
+    """
+    nc = 251
+    y_cases = [0, 1, 2, 58, 249, 250, 251, 2345]
+    n = len(y_cases)
+    distri_x = make_distri(nc, n)
+
+    graph = f"""
+        # some comments
+        NC {nc}
+        PROPERTY z = x+y+w
+        VAR MULTI z
+        VAR MULTI x
+        PUB MULTI y
+        PUB SINGLE w
+        """
+    y_cases = np.array(y_cases, dtype=np.uint32)
+    graph = FactorGraph(graph)
+    w = 5
+    for evidence_var, target_var, sub in [("x", "z", False), ("z", "x", True)]:
+        print("sub", sub)
+        bp_state = BPState(graph, n, {"y": y_cases, "w": w})
+        bp_state.set_evidence(evidence_var, distri_x)
+
+        bp_state.bp_loopy(1, True)
+        distri_z = bp_state.get_distribution(target_var)
+
+        distri_z_ref = np.zeros(distri_z.shape)
+
+        for x in range(nc):
+            for j in range(n):
+                if sub:
+                    new_idx = (x - y_cases[j] - w) % nc
+                else:
+                    new_idx = (x + y_cases[j] + w) % nc
+                distri_z_ref[j, new_idx] = distri_x[j, x]
+
+        assert np.allclose(distri_z_ref, distri_z)
+
+
+def test_ADD_multiple():
+    """
+    Test ADD between distributions
+    """
+    nc = 17
+    n = 4
+    distri_x = make_distri(nc, n)
+    distri_y = make_distri(nc, n)
+    distri_w = make_distri(nc, n)
+
+    graph = f"""
+        # some comments
+        NC {nc}
+        PROPERTY z = x+y+w
+        VAR MULTI z
+        VAR MULTI x
+        VAR MULTI y
+        VAR MULTI w
+        """
+
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n)
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+    bp_state.set_evidence("w", distri_w)
+
+    bp_state.bp_loopy(1, True)
+    distri_z = bp_state.get_distribution("z")
+
+    distri_z_ref = np.zeros(distri_z.shape)
+    msg = np.zeros(distri_z.shape)
+
+    for x in range(nc):
+        for y in range(nc):
+            for w in range(nc):
+                distri_z_ref[:, (x + y + w) % nc] += (
+                    distri_x[:, x] * distri_y[:, y] * distri_w[:, w]
+                )
+
+    distri_z_ref = (distri_z_ref.T / np.sum(distri_z_ref, axis=1)).T
+    assert np.allclose(distri_z_ref, distri_z)
+
+
+def test_MUL():
+    """
+    Test MUL between distributions
+    """
+    nc = 251
+    n = 4
+    distri_x = make_distri(nc, n)
+    distri_y = make_distri(nc, n)
+
+    graph = f"""
+        # some comments
+        NC {nc}
+        PROPERTY z = x*y
+        VAR MULTI z
+        VAR MULTI x
+        VAR MULTI y
+
+        """
+
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n)
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+
+    bp_state.bp_loopy(1, True)
+    distri_z = bp_state.get_distribution("z")
+
+    distri_z_ref = np.zeros(distri_z.shape)
+    msg = np.zeros(distri_z.shape)
+
+    for x in range(nc):
+        for y in range(nc):
+            distri_z_ref[:, (x * y) % nc] += distri_x[:, x] * distri_y[:, y]
+
+    distri_z_ref = (distri_z_ref.T / np.sum(distri_z_ref, axis=1)).T
+    assert np.allclose(distri_z_ref, distri_z)
+
+
+def test_xor():
+    """
+    Test XOR between distributions
+    """
+    nc = 512
+    n = 10
+    distri_x = np.random.randint(1, 10000000, (n, nc))
+    distri_x = (distri_x.T / np.sum(distri_x, axis=1)).T
+    distri_y = np.random.randint(1, 10000000, (n, nc))
+    distri_y = (distri_y.T / np.sum(distri_y, axis=1)).T
+
+    distri_b = np.random.randint(1, 10000000, (n, nc))
+    distri_b[1, 0] = 1.0
+    distri_b[1, 1] = 0.0
+    distri_b = (distri_b.T / np.sum(distri_b, axis=1)).T
+    distri_a = np.random.randint(1, 10000000, (n, nc))
+    distri_a[1, :] = 1.0
+    distri_a = (distri_a.T / np.sum(distri_a, axis=1)).T
+
+    graph = f"""
+        # some comments
+        NC {nc}
+        PROPERTY s2: z = x^y
+        PROPERTY s1: x = a^b
+        VAR MULTI x # some comments too
+        VAR MULTI y
+        VAR MULTI a
+        VAR MULTI b
+        VAR SINGLE z"""
+
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n)
+    bp_state.set_evidence("x", distri_x)
+    bp_state.set_evidence("y", distri_y)
+    bp_state.set_evidence("a", distri_a)
+    bp_state.set_evidence("b", distri_b)
+
+    bp_state2 = copy.deepcopy(bp_state)
+
+    # Custom optimized sequence -- equivelent to non-loopy algo, but focusing only on z.
+    bp_state.propagate_var("a")
+    bp_state.propagate_var("b")
+    bp_state.propagate_var("x")
+    bp_state.propagate_factor("s1")
+    bp_state.propagate_var("x")
+    bp_state.propagate_var("y")
+    bp_state.propagate_factor("s2")
+    bp_state.propagate_var("z")
+
+    bp_state2.bp_loopy(2, True)
+
+    distri_z = bp_state.get_distribution("z")
+    distri_z2 = bp_state2.get_distribution("z")
+
+    msg = np.zeros((n, nc))
+    distri_z_ref_multi = np.zeros((n, nc))
+    distri_z_ref = np.ones(distri_z.shape)
+
+    for a in range(nc):
+        for b in range(nc):
+            msg[:, a ^ b] += distri_a[:, a] * distri_b[:, b]
+
+    distri_x *= msg
+    for x in range(nc):
+        for y in range(nc):
+            distri_z_ref_multi[:, x ^ y] += distri_x[:, x] * distri_y[:, y]
+    for d in distri_z_ref_multi:
+        print(d)
+        distri_z_ref *= d
+
+    print("distri_x_ref", distri_x)
+    print("distri_y_ref", distri_y)
+    print("distri_z_ref", distri_z_ref)
+
+    distri_z_ref = distri_z_ref / np.sum(distri_z_ref)
+    assert np.allclose(distri_z_ref, distri_z)
+    assert np.allclose(distri_z_ref, distri_z2)
+
+
+def test_bad_var_norm():
+    """
+    When normalization of distributions can be bad...
+    Root cause of bug #55.
+    """
+    nc = 2
+    n = 1
+    dy = np.array([[1.0, 1e-100]])
+    dz = np.array([[1e-100, 1.0]])
+    graph = f"""
+        NC {nc}
+        VAR MULTI x
+        VAR MULTI y
+        VAR MULTI z
+        PUB MULTI a
+        PROPERTY s1: x = y^a
+        PROPERTY s2: x = !z
+        """
+
+    graph = FactorGraph(graph)
+    bp_state = BPState(graph, n, {"a": [0]})
+    bp_state.set_evidence("y", dy)
+    bp_state.set_evidence("z", dz)
+
+    bp_state.propagate_var("y")
+    bp_state.propagate_var("z")
+    bp_state.propagate_factor("s1")
+    bp_state.propagate_factor("s2")
+    bp_state.propagate_var("x")
+
+    s1 = normalize_distr(bp_state.get_belief_from_var("x", "s1"))
+    s2 = normalize_distr(bp_state.get_belief_from_var("x", "s2"))
+
+    assert np.allclose(np.array([[1.0, 1e-40]]), s1)
+    assert np.allclose(np.array([[1.0, 1e-40]]), s2)
+
+
+def test_cyclic():
+    graph_never_cyclic = """
+    NC 2
+    VAR MULTI x
+    VAR MULTI x0
+    VAR MULTI x1
+    VAR MULTI y
+    VAR MULTI y0
+    VAR MULTI y1
+    PUB MULTI p
+    VAR SINGLE k
+    VAR SINGLE k2
+
+    PROPERTY x = p ^ k
+    PROPERTY y = y0 ^ y1
+    PROPERTY x = x0 ^ x1
+    PROPERTY k2  = p ^ y
+    """
+    g = FactorGraph(graph_never_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() == False
+    graph_always_cyclic = """
+    NC 2
+    VAR MULTI x
+    VAR MULTI x0
+    VAR MULTI x1
+    VAR MULTI y
+    VAR MULTI y0
+    VAR MULTI y1
+    PUB MULTI p
+    VAR SINGLE k
+
+    PROPERTY x = p ^ k
+    PROPERTY y = y0 ^ y1
+    PROPERTY x = x0 ^ x1
+    PROPERTY y  = p ^ y1
+    """
+    g = FactorGraph(graph_always_cyclic)
+    assert BPState(g, 1, {"p": np.array([0], dtype=np.uint32)}).is_cyclic() == True
+    assert BPState(g, 2, {"p": np.array([0, 0], dtype=np.uint32)}).is_cyclic() == True
+    graph_multi_cyclic = """
+    NC 2
+    VAR MULTI x
+    VAR MULTI x0
+    VAR MULTI x1
+    VAR MULTI y
+    VAR MULTI y0
+    VAR MULTI y1
+    PUB MULTI p
+    VAR SINGLE k
+    VAR SINGLE k2
+
+    PROPERTY x = p ^ k
+    PROPERTY y = y0 ^ y1
+    PROPERTY x = x0 ^ x1
+    PROPERTY k2  = x ^ y
+    """
+    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
diff --git a/tests/test_sascagraph.py b/tests/test_sascagraph.py
index 751a5f2a..d4ac2194 100644
--- a/tests/test_sascagraph.py
+++ b/tests/test_sascagraph.py
@@ -23,6 +23,10 @@ def test_table():
             """
     graph = SASCAGraph(graph, n)
     graph.set_table("table", table)
+
+    x = distri_x.argmax(axis=1)
+    graph.sanity_check({"x": x, "y": table[x]})
+
     graph.set_init_distribution("x", distri_x)
 
     graph.run_bp(1)
@@ -36,6 +40,77 @@ def test_table():
     assert np.allclose(distri_y_ref, distri_y)
 
 
+def test_table_non_bij():
+    """
+    Test non-bijective Table lookup
+    """
+    nc = 2
+    n = 1
+    table = np.array([0, 0], dtype=np.uint32)
+    distri_x = np.array([[0.5, 0.5]])
+    distri_y = np.array([[0.8, 0.2]])
+
+    graph = f"""
+            PROPERTY y = table[x]
+            TABLE table
+            VAR MULTI x
+            VAR MULTI y
+            NC {nc}
+            """
+    graph = SASCAGraph(graph, n)
+    graph.set_table("table", table)
+
+    graph.sanity_check({"x": np.array([0, 1]), "y": np.array([0, 0])})
+
+    graph.set_init_distribution("x", distri_x)
+    graph.set_init_distribution("y", distri_y)
+
+    graph.run_bp(1)
+    distri_x = graph.get_distribution("x")
+    distri_y = graph.get_distribution("y")
+
+    distri_x_ref = np.array([0.5, 0.5])
+    distri_y_ref = np.array([1.0, 0.0])
+
+    assert np.allclose(distri_x_ref, distri_x)
+    assert np.allclose(distri_y_ref, distri_y)
+
+
+def test_not():
+    """
+    Test non-bijective Table lookup
+    """
+    nc = 2
+    n = 1
+    distri_x = np.array([[0.6, 0.4]])
+    distri_y = np.array([[0.8, 0.2]])
+
+    graph = f"""
+            PROPERTY y = !x
+            VAR MULTI x
+            VAR MULTI y
+            NC {nc}
+            """
+    graph = SASCAGraph(graph, n)
+
+    graph.sanity_check({"x": np.array([0, 1]), "y": np.array([1, 0])})
+
+    graph.set_init_distribution("x", distri_x)
+    graph.set_init_distribution("y", distri_y)
+
+    graph.run_bp(1)
+    distri_x = graph.get_distribution("x")
+    distri_y = graph.get_distribution("y")
+
+    s = 0.6 * 0.2 + 0.4 * 0.8
+    b = 0.6 * 0.2 / s
+    distri_x_ref = np.array([b, 1.0 - b])
+    distri_y_ref = np.array([1.0 - b, b])
+
+    assert np.allclose(distri_x_ref, distri_x)
+    assert np.allclose(distri_y_ref, distri_y)
+
+
 def test_and_public():
     """
     Test AND with public data
@@ -76,18 +151,21 @@ def test_xor_public():
     nc = 16
     n = 100
     public = np.random.randint(0, nc, n, dtype=np.uint32)
+    public2 = np.random.randint(0, nc, n, dtype=np.uint32)
     distri_x = np.random.randint(1, 100, (n, nc))
     distri_x = (distri_x.T / np.sum(distri_x, axis=1)).T
 
     graph = f"""
-        PROPERTY y = x ^ p
+        PROPERTY y = x ^ p ^ p2
         VAR MULTI y
         VAR MULTI x
         VAR MULTI p
+        VAR MULTI p2
         NC {nc}
         """
     graph = SASCAGraph(graph, n)
     graph.set_public("p", public)
+    graph.set_public("p2", public2)
     graph.set_init_distribution("x", distri_x)
 
     graph.run_bp(1)
@@ -95,7 +173,7 @@ def test_xor_public():
     distri_y = graph.get_distribution("y")
     distri_y_ref = np.zeros(distri_x.shape)
     for x in range(nc):
-        y = x ^ public
+        y = x ^ public ^ public2
         distri_y_ref[np.arange(n), y] = distri_x[np.arange(n), x]
 
     assert np.allclose(distri_y_ref, distri_y)
@@ -105,39 +183,97 @@ def test_AND():
     """
     Test AND between distributions
     """
-    nc = 256
-    n = 4
-    distri_x = np.random.randint(1, 10000000, (n, nc))
-    distri_x = (distri_x.T / np.sum(distri_x, axis=1)).T
-    distri_y = np.random.randint(1, 10000000, (n, nc))
-    distri_y = (distri_y.T / np.sum(distri_y, axis=1)).T
-
-    graph = f"""
-        # some comments
-        NC {nc}
-        PROPERTY z = x&y
-        VAR MULTI z
-        VAR MULTI x
-        VAR MULTI y
 
-        """
-
-    graph = SASCAGraph(graph, n)
-    graph.set_init_distribution("x", distri_x)
-    graph.set_init_distribution("y", distri_y)
-
-    graph.run_bp(1)
-    distri_z = graph.get_distribution("z")
-
-    distri_z_ref = np.zeros(distri_z.shape)
-    msg = np.zeros(distri_z.shape)
-
-    for x in range(nc):
-        for y in range(nc):
-            distri_z_ref[:, x & y] += distri_x[:, x] * distri_y[:, y]
+    def norm_d(d):
+        return d / np.sum(d, axis=1, keepdims=True)
+
+    def make_distri(nc, n):
+        return norm_d(np.random.randint(1, 10000000, (n, nc)).astype(np.float64))
+
+    cases = [
+        (
+            np.array([[0.5, 0.5]]),  # uniform x
+            np.array([[0.0, 1.0]]),  # y == 1
+            np.array([[0.0, 1.0]]),  # z == 1
+        ),
+        (
+            np.array([[0.1, 0.9]]),  # x == 1
+            np.array([[0.5, 0.5]]),  # uniform y
+            np.array([[1.0, 0.0]]),  # z == 1
+        ),
+        (
+            np.array([[0.1, 0.9]]),  # x == 1
+            np.array([[0.4, 0.6]]),  # uniform y
+            np.array([[0.8, 0.2]]),  # z == 1
+        ),
+        (make_distri(2, 1) for _ in range(3)),
+        (make_distri(4, 1) for _ in range(3)),
+        (make_distri(256, 4) for _ in range(3)),
+    ]
+
+    for distri_x, distri_y, distri_z in cases:
+        print("#### Test case:")
+        print(distri_x)
+        print(distri_y)
+        print(distri_z)
+        n, nc = distri_x.shape
+        graph = f"""
+            # some comments
+            NC {nc}
+            PROPERTY z = x&y
+            VAR MULTI z
+            VAR MULTI x
+            VAR MULTI y
 
-    distri_z_ref = (distri_z_ref.T / np.sum(distri_z_ref, axis=1)).T
-    assert np.allclose(distri_z_ref, distri_z)
+            """
+        graph = SASCAGraph(graph, n)
+        graph.set_init_distribution("x", distri_x)
+        graph.set_init_distribution("y", distri_y)
+        graph.set_init_distribution("z", distri_z)
+
+        distri_x_ref = np.zeros(distri_x.shape)
+        distri_y_ref = np.zeros(distri_y.shape)
+        distri_z_ref = np.zeros(distri_z.shape)
+
+        for x in range(nc):
+            for y in range(nc):
+                distri_x_ref[:, x] += distri_z[:, x & y] * distri_y[:, y]
+                distri_y_ref[:, y] += distri_z[:, x & y] * distri_x[:, x]
+                distri_z_ref[:, x & y] += distri_x[:, x] * distri_y[:, y]
+
+        # def make_dist_matrix(d):
+        #    d_matrix = np.zeros((n, nc, nc))
+        #    for x in range(nc):
+        #        for y in range(nc):
+        #            d_matrix[:,x&y,y] += d[:,x]
+        #    return d_matrix
+        # distri_x_matrix = make_dist_matrix(distri_x)
+        # distri_y_ref = np.linalg.solve(distri_x_matrix, distri_z)
+        # distri_y_matrix = make_dist_matrix(distri_y)
+        # distri_x_ref = np.linalg.solve(distri_y_matrix, distri_z)
+
+        distri_x_ref = norm_d(distri_x_ref * distri_x)
+        distri_y_ref = norm_d(distri_y_ref * distri_y)
+        distri_z_ref = norm_d(distri_z_ref * distri_z)
+
+        print("#### Ref:")
+        print(distri_x_ref)
+        print(distri_y_ref)
+        print(distri_z_ref)
+
+        graph.run_bp(1)
+        distri_x = graph.get_distribution("x")
+        distri_y = graph.get_distribution("y")
+        distri_z = graph.get_distribution("z")
+
+        print("#### Got:")
+        print(distri_x)
+        print(distri_y)
+        print(distri_z)
+
+        assert np.allclose(distri_z_ref, distri_z)
+        assert np.allclose(distri_x_ref, distri_x)
+        assert np.allclose(distri_y_ref, distri_y)
 
 
 def test_ADD():