--- a/trafficintelligence/prediction.py	Fri Apr 28 17:03:39 2023 -0400
+++ b/trafficintelligence/prediction.py	Tue May 02 07:12:26 2023 -0400
@@ -34,6 +34,11 @@
     def getPredictedSpeeds(self):
         return [so.norm for so in self.predictedSpeedOrientations.values()]
 
+    def hasCollisionWith(self, other, t, collisionDistanceThreshold):
+        p1 = self.predictPosition(t)
+        p2 = other.predictPosition(t)
+        return (p1-p2).norm2() <= collisionDistanceThreshold, p1, p2
+
     def plot(self, options = '', withOrigin = False, timeStep = 1, **kwargs):
         self.getPredictedTrajectory().plot(options, withOrigin, timeStep, **kwargs)
 
@@ -47,10 +52,35 @@
         self.probability = probability
         self.predictedPositions = {0: initialPosition}
         self.predictedSpeedOrientations = {0: moving.NormAngle.fromPoint(initialVelocity)}
-
+    
     def getControl(self):
         return self.control
 
+class PredictedTrajectoryConstantVelocity(PredictedTrajectory):
+    '''Predicted trajectory at constant speed or acceleration
+    TODO generalize by passing a series of velocities/accelerations'''
+
+    def __init__(self, initialPosition, initialVelocity, probability = 1.):
+        self.probability = probability
+        self.predictedPositions = {0: initialPosition}
+        self.initialVelocity = initialVelocity
+        #self.predictedSpeedOrientations = {0: moving.NormAngle.fromPoint(initialVelocity)}
+
+    def predictPosition(self, nTimeSteps):
+        if nTimeSteps > 0 and not nTimeSteps in self.predictedPositions:
+            self.predictPosition(nTimeSteps-1)
+            self.predictedPositions[nTimeSteps] = self.predictedPositions[nTimeSteps-1]+self.initialVelocity
+        return self.predictedPositions[nTimeSteps]
+
+class PredictedPolyTrajectoryConstant(PredictedTrajectory):
+    '''Predicted trajectory of an object with an outline represented by a polygon
+    Simple method that just translates the polygon corners'''
+    def __init__(self, polyCorners, initialVelocity, probability = 1.):
+        self.probability = probability
+        self.predictedPositions = {0: polyCorners}
+        self.initialVelocity = initialVelocity
+        self.predictedTrajectories = [PredictedTrajectoryConstant(p, initialVelocity) for p in polyCorners]
+        
 class PredictedTrajectoryPrototype(PredictedTrajectory):
     '''Predicted trajectory that follows a prototype trajectory
     The prototype is in the format of a moving.Trajectory: it could be
@@ -155,19 +185,15 @@
 
 def computeCollisionTime(predictedTrajectory1, predictedTrajectory2, collisionDistanceThreshold, timeHorizon):
     '''Computes the first instant 
-    at which two predicted trajectories are within some distance threshold
+    at which two predicted trajectories are colliding (depending on the trajectory format)
     Computes all the times including timeHorizon
     
     User has to check the first variable collision to know about a collision'''
     t = 1
-    p1 = predictedTrajectory1.predictPosition(t)
-    p2 = predictedTrajectory2.predictPosition(t)
-    collision = (p1-p2).norm2() <= collisionDistanceThreshold
+    collision, p1, p2 = predictedTrajectory1.hasCollisionWith(predictedTrajectory2, t, collisionDistanceThreshold)
     while t < timeHorizon and not collision:
         t += 1
-        p1 = predictedTrajectory1.predictPosition(t)
-        p2 = predictedTrajectory2.predictPosition(t)
-        collision = (p1-p2).norm2() <= collisionDistanceThreshold
+        collision, p1, p2 = predictedTrajectory1.hasCollisionWith(predictedTrajectory2, t, collisionDistanceThreshold)
     return collision, t, p1, p2
 
 def computeCollisionTimePolygon(predictedTrajectories1, predictedTrajectories2, timeHorizon):
@@ -185,7 +211,7 @@
         poly1 = moving.pointsToShapely([t1.predictPosition(t) for t1 in predictedTrajectories1])
         poly2 = moving.pointsToShapely([t2.predictPosition(t) for t2 in predictedTrajectories2])
         collision = poly1.overlaps(poly2)
-    return collision, t, p1, p2
+    return collision, t
 
 def savePredictedTrajectoriesFigure(currentInstant, obj1, obj2, predictedTrajectories1, predictedTrajectories2, timeHorizon, printFigure = True):
     from matplotlib.pyplot import figure, axis, title, clf, savefig
@@ -380,11 +406,11 @@
         return collisionProbabilities
 
 class ConstantPredictionParameters(PredictionParameters):
-    def __init__(self, maxSpeed):
-        PredictionParameters.__init__(self, 'constant velocity', maxSpeed)
+    def __init__(self):
+        PredictionParameters.__init__(self, 'constant velocity')
 
     def generatePredictedTrajectories(self, obj, instant):
-        return [PredictedTrajectoryConstant(obj.getPositionAtInstant(instant), obj.getVelocityAtInstant(instant), maxSpeed = self.maxSpeed)]
+        return [PredictedTrajectoryConstantVelocity(obj.getPositionAtInstant(instant), obj.getVelocityAtInstant(instant))]
 
 class NormalAdaptationPredictionParameters(PredictionParameters):
     def __init__(self, maxSpeed, nPredictedTrajectories, accelerationDistribution, steeringDistribution, useFeatures = False):
@@ -438,7 +464,7 @@
             velocities = [f.getVelocityAtInstant(instant) for f in features]
             probability = 1./float(len(positions))
             for initialPosition,initialVelocity in zip(positions, velocities):
-                predictedTrajectories.append(PredictedTrajectoryConstant(initialPosition, initialVelocity, probability = probability, maxSpeed = self.maxSpeed))
+                predictedTrajectories.append(PredictedTrajectoryConstantVelocity(initialPosition, initialVelocity, probability = probability))
             return predictedTrajectories
         else:
             print('Object {} has no features'.format(obj.getNum()))

--- a/trafficintelligence/tests/prediction.txt	Fri Apr 28 17:03:39 2023 -0400
+++ b/trafficintelligence/tests/prediction.txt	Tue May 02 07:12:26 2023 -0400
@@ -16,6 +16,10 @@
 >>> et.predictPosition(12) # doctest:+ELLIPSIS
 (19.5...,0.0...)
 
+>>> et = PredictedTrajectoryConstantVelocity(moving.Point(0,0), moving.Point(1,0))
+>>> et.predictPosition(10) # doctest:+ELLIPSIS
+(10.0...,0.0...)
+
 >>> import random
 >>> acceleration = lambda: random.uniform(-0.5,0.5)
 >>> steering = lambda: random.uniform(-0.1,0.1)