add ability to distribute next (push to bottom of AST)

stl/test_utils.py

@@ -86,6 +86,16 @@ def test_discretize():
     phi = stl.parse('G[0.3, 1.2]((AP1) U (AP2))')
     assert not stl.utils.is_discretizable(phi, dt)
 
+    phi = stl.parse('G[0.3, 0.6](~(F[0, 0.3](A)))')
+    assert stl.utils.is_discretizable(phi, dt)
+    phi2 = stl.utils.discretize(phi, dt, distribute=True)
+    phi3 = stl.utils.discretize(phi2, dt, distribute=True)
+    assert phi2 == phi3
+    assert phi2 == stl.parse(
+        '(~((X(A)) ∨ (X(X(A))))) ∧ (~((X(X(A))) ∨ (X(X(X(A))))))')
+
+
+
 
 def test_scope():
     dt = 0.3

stl/utils.py

@@ -122,8 +122,9 @@ def scope(phi, dt, *, _t=0):
 
 
 @require_discretizable
-def discretize(phi, dt):
-    return _discretize(phi, dt)
+def discretize(phi, dt, distribute=False):
+    phi = _discretize(phi, dt)
+    return _distribute_next(phi) if distribute else phi
 
 
 def _discretize(phi, dt):
@@ -151,6 +152,20 @@ def _interval_discretizable(itvl, dt):
     return np.isclose(l, round(l)) and np.isclose(u, round(u))
 
 
+def _distribute_next(phi, i=0):
+    if isinstance(phi, (LinEq, AtomicPred)):
+        return stl.utils.next(phi, i=i)
+    elif isinstance(phi, Next):
+        return _distribute_next(phi.arg, i=i+1)
+
+    children = tuple(_distribute_next(c, i) for c in phi.children)
+
+    if isinstance(phi, (And, Or)):
+        return bind(phi).args.set(children)
+    elif isinstance(phi, (Neg, Next)):
+        return bind(phi).arg.set(children[0])
+
+
 def is_discretizable(phi, dt):
     if any(c for c in phi.walk() if isinstance(c, Until)):
         return False