1 files changed, 68 insertions, 0 deletions
diff --git a/phstuff/barcode.py b/phstuff/barcode.py
new file mode 100644
index 0000000..9c7b05e
--- /dev/null
+++ b/phstuff/barcode.py
@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+class Interval:
+    def __init__(self, birth, death = None):
+        if death is not None and death < birth:
+            raise ValueError("Death must be at least birth.")
+        self.birth = birth
+        self.death = death
+
+    def is_finite(self):
+        return self.death is not None
+    
+    def __str__(self):
+        if self.is_finite():
+            return "(%s, %s)" %(str(self.birth), str(self.death))
+        else:
+            return "(%s, ∞)" %(str(self.birth))
+
+    def __repr__(self):
+        if self.is_finite():
+            return "Interval(%s, %s)" %(repr(self.birth), repr(self.death))
+        else:
+            return "Interval(%s)" %(repr(self.birth))
+
+    def size(self):
+        if self.is_finite():
+            return self.death - self.birth
+        else:
+            return float("inf")
+
+
+def finitize(barcode, max_scale):
+    ret = []
+    for bar in barcode:
+        if bar.is_finite():
+            ret.append(bar)
+        else:
+            ret.append(Interval(bar.birth, max_scale))
+    return ret
+    
+def betti_curve(barcode, min_scale, max_scale):
+    xs = []
+    for interval in pd:
+        xs.append((interval.l, 1))
+        if interval.is_finite():
+            xs.append((interval.u, -1))
+
+    if not xs:
+        return np.zeros((2,1))
+    
+    xs.sort()
+    scales = [min_scale]
+    bettis = [0]
+    
+    for (t, delta) in xs:
+        if scales[-1] == t:
+            bettis[-1] += delta
+        else:
+            scales.append(t)
+            bettis.append(bettis[-1] + delta)
+
+    if scales[-1] < max_scale:
+        scales.append(max_scale)
+        bettis.append(bettis[-1])
+        
+    return (np.array(scales), np.array(bettis, dtype=int))