Use external dataset in scheduler

NASA-AMMOS · Sep 15, 2023 · 8ff37f3 · 8ff37f3
1 parent 6a77a90
commit 8ff37f3
Show file tree

Hide file tree

Showing 10 changed files with 183 additions and 35 deletions.
diff --git a/constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/SpansWrapperExpression.java b/constraints/src/main/java/gov/nasa/jpl/aerie/constraints/tree/SpansWrapperExpression.java
@@ -0,0 +1,29 @@
+package gov.nasa.jpl.aerie.constraints.tree;
+import gov.nasa.jpl.aerie.constraints.model.EvaluationEnvironment;
+import gov.nasa.jpl.aerie.constraints.model.SimulationResults;
+import gov.nasa.jpl.aerie.constraints.time.Interval;
+import gov.nasa.jpl.aerie.constraints.time.Spans;
+
+import java.util.Set;
+
+public record SpansWrapperExpression(Spans spans) implements Expression<Spans> {
+  @Override
+  public Spans evaluate(
+      final SimulationResults results,
+      final Interval bounds,
+      final EvaluationEnvironment environment)
+  {
+    return spans;
+  }
+
+  @Override
+  public String prettyPrint(final String prefix) {
+    return String.format(
+        "\n%s(spans-wrapper-of %s)",
+        prefix,
+        this.spans
+    );  }
+
+  @Override
+  public void extractResources(final Set<String> names) { }
+}
diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/CardinalityGoal.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/CardinalityGoal.java
@@ -121,10 +121,10 @@ protected CardinalityGoal fill(CardinalityGoal goal) {
    * should probably be created!)
    */
   @Override
-  public Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults) {
+  public Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults, final EvaluationEnvironment evaluationEnvironment) {
 
     //unwrap temporalContext
-    final var windows = getTemporalContext().evaluate(simulationResults);
+    final var windows = getTemporalContext().evaluate(simulationResults, evaluationEnvironment);
 
     //make sure it hasn't changed
     if (this.initiallyEvaluatedTemporalContext != null && !windows.equals(this.initiallyEvaluatedTemporalContext)) {
@@ -142,7 +142,7 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
       final var actTB =
           new ActivityExpression.Builder().basedOn(this.matchActTemplate).startsOrEndsIn(subIntervalWindows).build();
 
-      final var acts = new LinkedList<>(plan.find(actTB, simulationResults, new EvaluationEnvironment()));
+      final var acts = new LinkedList<>(plan.find(actTB, simulationResults, evaluationEnvironment));
       acts.sort(Comparator.comparing(SchedulingActivityDirective::startOffset));
 
       int nbActs = 0;
@@ -201,7 +201,7 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
             this,
             subIntervalWindows,
             this.desiredActTemplate,
-            new EvaluationEnvironment(),
+            evaluationEnvironment,
             nbToSchedule,
             durToSchedule.isPositive() ? Optional.of(durToSchedule) : Optional.empty()));
       }

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/CoexistenceGoal.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/CoexistenceGoal.java
@@ -19,7 +19,7 @@
 import gov.nasa.jpl.aerie.scheduler.model.SchedulingActivityDirective;
 
 import java.util.ArrayList;
-import java.util.Map;
+import java.util.HashMap;
 import java.util.Optional;
 
 /**
@@ -172,7 +172,8 @@ protected CoexistenceGoal fill(CoexistenceGoal goal) {
    * should probably be created!)
    */
   @SuppressWarnings({"unchecked", "rawtypes"})
-  public java.util.Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults) { //TODO: check if interval gets split and if so, notify user?
+  @Override
+  public java.util.Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults, final EvaluationEnvironment evaluationEnvironment) { //TODO: check if interval gets split and if so, notify user?
 
     //NOTE: temporalContext IS A WINDOWS OVER WHICH THE GOAL APPLIES, USUALLY SOMETHING BROAD LIKE A MISSION PHASE
     //NOTE: expr IS A WINDOWS OVER WHICH A COEXISTENCEGOAL APPLIES, FOR EXAMPLE THE WINDOWS CORRESPONDING TO 5 SECONDS AFTER EVERY BASICACTIVITY IS SCHEDULED
@@ -182,7 +183,7 @@ public java.util.Collection<Conflict> getConflicts(Plan plan, final SimulationRe
     //        AN ACTIVITYEXPRESSION AND THEN ANALYZEWHEN WAS A MISSION PHASE, ALTHOUGH IT IS POSSIBLE TO JUST SPECIFY AN EXPRESSION<WINDOWS> THAT COMBINES THOSE.
 
     //unwrap temporalContext
-    final var windows = getTemporalContext().evaluate(simulationResults);
+    final var windows = getTemporalContext().evaluate(simulationResults, evaluationEnvironment);
 
     //make sure it hasn't changed
     if (this.initiallyEvaluatedTemporalContext != null && !windows.includes(this.initiallyEvaluatedTemporalContext)) {
@@ -192,7 +193,7 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
       this.initiallyEvaluatedTemporalContext = windows;
     }
 
-    final var anchors = expr.evaluate(simulationResults).intersectWith(windows);
+    final var anchors = expr.evaluate(simulationResults, evaluationEnvironment).intersectWith(windows);
 
     //make sure expr hasn't changed either as that could yield unexpected behavior
     if (this.evaluatedExpr != null && !anchors.isCollectionSubsetOf(this.evaluatedExpr)) {
@@ -242,7 +243,10 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
         activityCreationTemplate.durationIn(durRange);
       }
 
-      final var existingActs = plan.find(activityFinder.build(), simulationResults, createEvaluationEnvironmentFromAnchor(window));
+      final var existingActs = plan.find(
+          activityFinder.build(),
+          simulationResults,
+          createEvaluationEnvironmentFromAnchor(evaluationEnvironment, window));
 
       var missingActAssociations = new ArrayList<SchedulingActivityDirective>();
       var planEvaluation = plan.getEvaluation();
@@ -272,7 +276,7 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
       if (!alreadyOneActivityAssociated) {
         //create conflict if no matching target activity found
         if (existingActs.isEmpty()) {
-          conflicts.add(new MissingActivityTemplateConflict(this, this.temporalContext.evaluate(simulationResults), temp, createEvaluationEnvironmentFromAnchor(window), 1, Optional.empty()));
+          conflicts.add(new MissingActivityTemplateConflict(this, this.temporalContext.evaluate(simulationResults, evaluationEnvironment), temp, createEvaluationEnvironmentFromAnchor(evaluationEnvironment, window), 1, Optional.empty()));
         } else {
           conflicts.add(new MissingAssociationConflict(this, missingActAssociations));
         }
@@ -283,24 +287,28 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
     return conflicts;
   }
 
-  private EvaluationEnvironment createEvaluationEnvironmentFromAnchor(Segment<Optional<Spans.Metadata>> span){
+  private EvaluationEnvironment createEvaluationEnvironmentFromAnchor(EvaluationEnvironment existingEnvironment, Segment<Optional<Spans.Metadata>> span){
     if(span.value().isPresent()){
       final var metadata = span.value().get();
+      final var activityInstances = new HashMap<>(existingEnvironment.activityInstances());
+      activityInstances.put(this.alias, metadata.activityInstance());
       return new EvaluationEnvironment(
-          Map.of(this.alias, metadata.activityInstance()),
-          Map.of(),
-          Map.of(),
-          Map.of(),
-          Map.of()
+          activityInstances,
+          existingEnvironment.spansInstances(),
+          existingEnvironment.intervals(),
+          existingEnvironment.realExternalProfiles(),
+          existingEnvironment.discreteExternalProfiles()
       );
     } else{
       assert this.alias != null;
+      final var intervals = new HashMap<>(existingEnvironment.intervals());
+      intervals.put(this.alias, span.interval());
       return new EvaluationEnvironment(
-          Map.of(),
-          Map.of(),
-          Map.of(this.alias, span.interval()),
-          Map.of(),
-          Map.of()
+          existingEnvironment.activityInstances(),
+          existingEnvironment.spansInstances(),
+          intervals,
+          existingEnvironment.realExternalProfiles(),
+          existingEnvironment.discreteExternalProfiles()
       );
     }
   }

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/Goal.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/Goal.java
@@ -1,5 +1,6 @@
 package gov.nasa.jpl.aerie.scheduler.goals;
 
+import gov.nasa.jpl.aerie.constraints.model.EvaluationEnvironment;
 import gov.nasa.jpl.aerie.constraints.model.SimulationResults;
 import gov.nasa.jpl.aerie.constraints.time.Interval;
 import gov.nasa.jpl.aerie.constraints.time.Windows;
@@ -283,7 +284,8 @@ public String getName() {
    * @param simulationResults
    * @return a list of issues in the plan that diminish goal satisfaction
    */
-  public java.util.Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults) {
+  public java.util.Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults, final
+                                                     EvaluationEnvironment evaluationEnvironment) {
     return java.util.Collections.emptyList();
   }
 

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/OptionGoal.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/OptionGoal.java
@@ -1,5 +1,6 @@
 package gov.nasa.jpl.aerie.scheduler.goals;
 
+import gov.nasa.jpl.aerie.constraints.model.EvaluationEnvironment;
 import gov.nasa.jpl.aerie.constraints.model.SimulationResults;
 import gov.nasa.jpl.aerie.scheduler.conflicts.Conflict;
 import gov.nasa.jpl.aerie.scheduler.model.Plan;
@@ -28,7 +29,9 @@ public Optimizer getOptimizer(){
   }
 
   @Override
-  public java.util.Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults) {
+  public java.util.Collection<Conflict> getConflicts(Plan plan,
+                                                     final SimulationResults simulationResults,
+                                                     final EvaluationEnvironment evaluationEnvironment) {
     throw new NotImplementedException("Conflict detection is performed at solver level");
   }
 

diff --git a/...duler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/ProceduralCreationGoal.java b/...duler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/ProceduralCreationGoal.java
@@ -108,12 +108,13 @@ protected ProceduralCreationGoal fill(ProceduralCreationGoal goal) {
    * the plan (and should probably be created). The matching is strict: all
    * arguments must be identical.
    */
-  public Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults) {
+  @Override
+  public Collection<Conflict> getConflicts(Plan plan, final SimulationResults simulationResults, final EvaluationEnvironment evaluationEnvironment) {
     final var conflicts = new java.util.LinkedList<Conflict>();
 
     //run the generator to see what acts are still desired
     //REVIEW: maybe some caching on plan hash here?
-    final var requestedActs = getRelevantGeneratedActivities(plan, simulationResults);
+    final var requestedActs = getRelevantGeneratedActivities(plan, simulationResults, evaluationEnvironment);
 
     //walk each requested act and try to find an exact match in the plan
     for (final var requestedAct : requestedActs) {
@@ -202,13 +203,13 @@ protected ProceduralCreationGoal() { }
    *     are deemed relevant to this goal (eg within the temporal context
    *     of this goal)
    */
-  private Collection<SchedulingActivityDirective> getRelevantGeneratedActivities(Plan plan, SimulationResults simulationResults) {
+  private Collection<SchedulingActivityDirective> getRelevantGeneratedActivities(Plan plan, SimulationResults simulationResults, EvaluationEnvironment evaluationEnvironment) {
 
     //run the generator in the plan context
     final var allActs = generator.apply(plan);
 
     //filter out acts that don't have a start time within the goal purview
-    final var evaluatedGoalContext = getTemporalContext().evaluate(simulationResults);
+    final var evaluatedGoalContext = getTemporalContext().evaluate(simulationResults, evaluationEnvironment);
     final var filteredActs = allActs.stream().filter(
         act -> ((act.startOffset() != null)
                 && evaluatedGoalContext.includes(Interval.at(0, act.startOffset())))

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/RecurrenceGoal.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/goals/RecurrenceGoal.java
@@ -104,12 +104,13 @@ else if (every.max.isNegative()) {
    * exist over a timespan longer than the allowed range (and one should
    * probably be created!)
    */
-  public java.util.Collection<Conflict> getConflicts(@NotNull Plan plan, final SimulationResults simulationResults) {
+  @Override
+  public java.util.Collection<Conflict> getConflicts(@NotNull Plan plan, final SimulationResults simulationResults, final EvaluationEnvironment evaluationEnvironment) {
     final var conflicts = new java.util.LinkedList<Conflict>();
 
     //unwrap temporalContext
     final var tempWindowPlanHorizon = new Windows(false).set(List.of(this.planHorizon.getHor()), true);
-    final var windows = tempWindowPlanHorizon.and(this.getTemporalContext().evaluate(simulationResults));
+    final var windows = tempWindowPlanHorizon.and(this.getTemporalContext().evaluate(simulationResults, evaluationEnvironment));
 
     //check repeat is larger than activity duration
     if(this.getActTemplate().getType().getDurationType() instanceof DurationType.Fixed act){
@@ -153,7 +154,7 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
         final var strideDur = actStartT.minus(prevStartT);
         if (strideDur.compareTo(this.recurrenceInterval.max) > 0) {
           //fill conflicts for all the missing activities in that long span
-          conflicts.addAll(makeRecurrenceConflicts(prevStartT, actStartT));
+          conflicts.addAll(makeRecurrenceConflicts(prevStartT, actStartT, evaluationEnvironment));
 
         } else {
           /*TODO: right now, we associate with all the activities that are satisfying but we should aim for the minimum
@@ -171,7 +172,7 @@ else if (this.initiallyEvaluatedTemporalContext == null) {
 
       //fill in conflicts for all missing activities in the last span up to the
       //goal's own end time (also handles case of no matching acts at all)
-      conflicts.addAll(makeRecurrenceConflicts(prevStartT, lastStartT));
+      conflicts.addAll(makeRecurrenceConflicts(prevStartT, lastStartT, evaluationEnvironment));
     }
 
     return conflicts;
@@ -202,7 +203,7 @@ protected RecurrenceGoal() { }
    * @param start IN the start time of the span to fill with conflicts (inclusive)
    * @param end IN the end time of the span to fill with conflicts (exclusive)
    */
-  private java.util.Collection<MissingActivityConflict> makeRecurrenceConflicts(Duration start, Duration end)
+  private java.util.Collection<MissingActivityConflict> makeRecurrenceConflicts(Duration start, Duration end, final EvaluationEnvironment evaluationEnvironment)
   {
     final var conflicts = new java.util.LinkedList<MissingActivityConflict>();
 
@@ -212,7 +213,7 @@ private java.util.Collection<MissingActivityConflict> makeRecurrenceConflicts(Du
     ) {
       final var windows = new Windows(false).set(Interval.betweenClosedOpen(intervalT.minus(recurrenceInterval.max), Duration.min(intervalT, end)), true);
       if(windows.iterateEqualTo(true).iterator().hasNext()){
-        conflicts.add(new MissingActivityTemplateConflict(this, windows, this.getActTemplate(), new EvaluationEnvironment(), 1, Optional.empty()));
+        conflicts.add(new MissingActivityTemplateConflict(this, windows, this.getActTemplate(), evaluationEnvironment, 1, Optional.empty()));
       }
       else{
         System.out.println();

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/model/Problem.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/model/Problem.java
@@ -1,5 +1,7 @@
 package gov.nasa.jpl.aerie.scheduler.model;
 
+import gov.nasa.jpl.aerie.constraints.model.DiscreteProfile;
+import gov.nasa.jpl.aerie.constraints.model.LinearProfile;
 import gov.nasa.jpl.aerie.merlin.driver.MissionModel;
 import gov.nasa.jpl.aerie.merlin.driver.SimulationResults;
 import gov.nasa.jpl.aerie.merlin.protocol.model.SchedulerModel;
@@ -12,7 +14,9 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -40,6 +44,10 @@ public class Problem {
    */
   private final List<GlobalConstraint> globalConstraints
       = new java.util.LinkedList<>();
+
+  private final Map<String, LinearProfile> realExternalProfiles = new HashMap<>();
+  private final Map<String, DiscreteProfile> discreteExternalProfiles = new HashMap<>();
+
   /**
    * the initial seed plan to start scheduling from
    */
@@ -148,6 +156,20 @@ public void setInitialPlan(final Plan plan) {
 
   public Optional<SimulationData> getInitialSimulationResults(){ return initialSimulationResults; }
 
+  public void setExternalProfile(final Map<String, LinearProfile> realExternalProfiles,
+                                  final Map<String, DiscreteProfile> discreteExternalProfiles){
+    this.realExternalProfiles.putAll(realExternalProfiles);
+    this.discreteExternalProfiles.putAll(discreteExternalProfiles);
+  }
+
+  public Map<String, LinearProfile> getRealExternalProfiles(){
+    return this.realExternalProfiles;
+  }
+
+  public Map<String, DiscreteProfile> getDiscreteExternalProfiles(){
+    return this.discreteExternalProfiles;
+  }
+
   public void setGoals(List<Goal> goals){
     goalsOrderedByPriority.clear();
     goalsOrderedByPriority.addAll(goals);

diff --git a/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/solver/PrioritySolver.java b/scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/solver/PrioritySolver.java
@@ -613,7 +613,8 @@ else if(!analysisOnly &&  (missing instanceof MissingActivityTemplateConflict mi
     assert plan != null;
     //REVIEW: maybe should have way to request only certain kinds of conflicts
     final var lastSimulationResults = this.getLatestSimResultsUpTo(this.problem.getPlanningHorizon().getEndAerie());
-    final var rawConflicts = goal.getConflicts(plan, lastSimulationResults);
+    final var evaluationEnvironment = new EvaluationEnvironment(this.problem.getRealExternalProfiles(), this.problem.getDiscreteExternalProfiles());
+    final var rawConflicts = goal.getConflicts(plan, lastSimulationResults, evaluationEnvironment);
     assert rawConflicts != null;
     return rawConflicts;
   }
@@ -765,9 +766,10 @@ private Windows narrowByResourceConstraints(Windows windows,
     final var latestSimulationResults = this.getLatestSimResultsUpTo(totalDomain.end);
     //iteratively narrow the windows from each constraint
     //REVIEW: could be some optimization in constraint ordering (smallest domain first to fail fast)
+    final var evaluationEnvironment = new EvaluationEnvironment(this.problem.getRealExternalProfiles(), this.problem.getDiscreteExternalProfiles());
     for (final var constraint : constraints) {
       //REVIEW: loop through windows more efficient than enveloppe(windows) ?
-      final var validity = constraint.evaluate(latestSimulationResults, totalDomain);
+      final var validity = constraint.evaluate(latestSimulationResults, totalDomain, evaluationEnvironment);
       ret = ret.and(validity);
       //short-circuit if no possible windows left
       if (ret.stream().noneMatch(Segment::value)) {