adding transfer entropy

Author: Mattehub

hoi/core/combinatory.py

@@ -1,12 +1,13 @@
 import itertools
 from math import comb as ccomb
+from math import perm as pperm
 
 import jax.numpy as jnp
 import numpy as np
 
 
-def _combinations(n, k, order, target):
-    for c in itertools.combinations(range(n), k):
+def _combinations(n, k, order, target, fnc=None):
+    for c in fnc(range(n), k):
         # convert to list
         c = list(c) + target
 
@@ -25,6 +26,7 @@ def combinations(
     order=False,
     fill_value=-1,
     target=None,
+    directed=False,
 ):
     """Get combinations.
 
@@ -53,13 +55,21 @@ def combinations(
         combinations of k elements.
     """
     # ________________________________ ITERATOR _______________________________
+
+    if directed:
+        fnc = itertools.permutations
+        fnc_nmult = pperm
+    else:
+        fnc = itertools.combinations
+        fnc_nmult = ccomb
+
     if not isinstance(maxsize, int):
         maxsize = minsize
     target = [] if target is None else list(target)
     assert maxsize >= minsize
     iterators = []
     for msize in range(minsize, maxsize + 1):
-        iterators.append(_combinations(n, msize, order, target))
+        iterators.append(_combinations(n, msize, order, target, fnc))
     iterators = itertools.chain(*tuple(iterators))
 
     if astype == "iterator":
@@ -70,7 +80,7 @@ def combinations(
         combs = np.asarray([c for c in iterators]).astype(int)
     else:
         # get the number of combinations
-        n_mults = sum([ccomb(n, c) for c in range(minsize, maxsize + 1)])
+        n_mults = sum([fnc_nmult(n, c) for c in range(minsize, maxsize + 1)])
 
         # prepare output
         combs = np.full(

hoi/core/redundancies.py

@@ -12,3 +12,4 @@
 from .syn_mmi import SynergyMMI  # noqa
 from .tc import TC  # noqa
 from .pairwise_mi import MI  # noqa
+from .transfer_entropy import TransferEntropy  # noqa

hoi/metrics/__init__.py

@@ -283,6 +283,7 @@ def get_combinations(self, minsize, maxsize=None, astype="jax"):
                 astype=astype,
                 order=False,
                 fill_value=-1,
+                directed=self._directed,
             )
 
         return self._multiplets, self.order

hoi/metrics/base_hoi.py

@@ -72,6 +72,7 @@ class DOtot(HOIEstimator):
     _positive = "redundancy"
     _negative = "synergy"
     _symmetric = False
+    _directed = False
 
     def __init__(self, x, multiplets=None, verbose=None):
         HOIEstimator.__init__(

hoi/metrics/do_tot.py

@@ -66,6 +66,7 @@ class DTC(HOIEstimator):
     _negative = "null"
     _positive = "null"
     _symmetric = True
+    _directed = False
 
     def __init__(self, x, y=None, multiplets=None, verbose=None):
         HOIEstimator.__init__(

hoi/metrics/dtc.py

@@ -42,6 +42,7 @@ class GradientOinfo(HOIEstimator):
     _positive = "redundancy"
     _negative = "synergy"
     _symmetric = True
+    _directed = False
 
     def __init__(self, x, y, multiplets=None, base_model=Oinfo, verbose=None):
         kw_oinfo = dict(multiplets=multiplets, verbose=verbose)

hoi/metrics/gradient_oinfo.py

@@ -38,6 +38,7 @@ class InfoTot(HOIEstimator):
     _positive = "info"
     _negative = "null"
     _symmetric = False
+    _directed = False
 
     def __init__(self, x, y, multiplets=None, verbose=None):
         HOIEstimator.__init__(

hoi/metrics/info_tot.py

@@ -75,6 +75,7 @@ class InfoTopo(HOIEstimator):
     _positive = "redundancy"
     _negative = "synergy"
     _symmetric = True
+    _directed = False
 
     def __init__(self, x, y=None, verbose=None):
         # for infotopo, the multiplets are set to None because this metric

hoi/metrics/infotopo.py

@@ -71,6 +71,7 @@ class Oinfo(HOIEstimator):
     _positive = "redundancy"
     _negative = "synergy"
     _symmetric = True
+    _directed = False
 
     def __init__(self, x, y=None, multiplets=None, verbose=None):
         HOIEstimator.__init__(