created a weighted scorer and started writing tests

astroplan/constraints.py

@@ -529,7 +529,7 @@ class MoonIlluminationConstraint(Constraint):
 
     Constraint is also satisfied if the Moon has set.
     """
-    def __init__(self, min=None, max=None, ephemeris=None):
+    def __init__(self, min=None, max=None, ephemeris=None, boolean_constraint=True):
         """
         Parameters
         ----------
@@ -543,10 +543,15 @@ def __init__(self, min=None, max=None, ephemeris=None):
             Ephemeris to use.  If not given, use the one set with
             `~astropy.coordinates.solar_system_ephemeris` (which is
             set to 'builtin' by default).
+        boolean_constraint : bool
+            If True, the constraint is treated as a boolean (True for within the
+            limits and False for outside).  If False, the constraint returns a
+            float on [0, 1], where 0 is the min altitude and 1 is the max.
         """
         self.min = min
         self.max = max
         self.ephemeris = ephemeris
+        self.boolean_constraint = boolean_constraint
 
     @classmethod
     def dark(cls, min=None, max=0.25, **kwargs):
@@ -607,16 +612,40 @@ def compute_constraint(self, times, observer, targets):
         moon_up_mask = moon_alt >= 0
 
         illumination = cached_moon['illum']
-        if self.min is None and self.max is not None:
-            mask = (self.max >= illumination) | moon_down_mask
-        elif self.max is None and self.min is not None:
-            mask = (self.min <= illumination) & moon_up_mask
-        elif self.min is not None and self.max is not None:
-            mask = ((self.min <= illumination) &
-                    (illumination <= self.max)) & moon_up_mask
+        if self.boolean_constraint:
+            if self.min is None and self.max is not None:
+                mask = (self.max >= illumination) | moon_down_mask
+            elif self.max is None and self.min is not None:
+                mask = (self.min <= illumination) & moon_up_mask
+            elif self.min is not None and self.max is not None:
+                mask = ((self.min <= illumination) &
+                        (illumination <= self.max)) & moon_up_mask
+            else:
+                raise ValueError("No max and/or min specified in "
+                                 "MoonSeparationConstraint.")
         else:
-            raise ValueError("No max and/or min specified in "
-                             "MoonSeparationConstraint.")
+            if self.min is None and self.max is not None:
+                moon_down = np.where(moon_down_mask == 1)
+                mask = min_best_rescale(illumination, 0, self.max, 0)
+                mask[moon_down] = 1
+            elif self.max is None and self.min is not None:
+                moon_down = np.where(moon_down_mask == 1)
+                mask = min_best_rescale(illumination, self.min, 1, 0)
+                if self.min == 0:
+                    mask[moon_down] = 1
+                else:
+                    mask[moon_down] = 0
+            elif self.min is not None and self.max is not None:
+                moon_down = np.where(moon_down_mask == 1)
+                mask = min_best_rescale(illumination, self.min,
+                                        self.max, 0)
+                if self.min == 0:
+                    mask[moon_down] = 1
+                else:
+                    mask[moon_down] = 0
+            else:
+                raise ValueError("No max and/or min specified in "
+                                 "MoonSeparationConstraint.")
 
         if targets is not None:
             mask = np.tile(mask, len(targets))

astroplan/scheduling.py

@@ -147,6 +147,83 @@ def create_score_array(self, time_resolution=1*u.minute):
             score_array *= constraint(self.observer, targets, times)
         return score_array
 
+    def weighted_score_array(self, time_resolution=1*u.minute):
+        """
+        For each block, this returns a score for each time using the
+        formula (score*weight+score*weight+...)/(weight+weight+..)
+        """
+        # note: with this method, a non-boolean constraint with weight
+        # zero has the exact same effect if it was boolean
+        start = self.schedule.start_time
+        end = self.schedule.end_time
+        times = time_grid_from_range((start, end), time_resolution)
+        # create an array to hold all of the scores
+        score_array = np.zeros((len(self.blocks), len(times)))
+        # create an array to record where any of the constraints are zero
+        constraint_zeros = np.ones((len(self.blocks), len(times)), dtype=int)
+        local_constraints = []
+        weights = []
+        for i, block in enumerate(self.blocks):
+            weights.append(0)
+            local_constraints.append([])
+            # schedulers put global constraints into ._all_constraints
+            # so we can use .constraints for the local constraints
+            if block.constraints:
+                for constraint in block.constraints:
+                    local_constraints[i].append(constraint.__class__.__name__)
+                    applied_score = constraint(self.observer, [block.target],
+                                               times=times)[0]
+                    if constraint.boolean_constraint:
+                        # add to the mask designating where the score is zero
+                        # if either is 0, constraint_zeros becomes 0
+                        constraint_zeros[i] &= applied_score
+                    # TODO: make a default weight=1 and merge these
+                    elif constraint.weight:
+                        constraint_zeros[i][(applied_score == 0)] = 0
+                        weight = constraint.weight
+                        weights[i] += weight
+                        score_array[i] += applied_score * weight
+                    else:
+                        constraint_zeros[i][(applied_score == 0)] = 0
+                        weights[i] += 1
+                        score_array[i] += applied_score
+        targets = [block.target for block in self.blocks]
+        for constraint in self.global_constraints:
+            skip_global = []
+            for i, block in enumerate(local_constraints):
+                if constraint.__class__.__name__ in block:
+                    skip_global.append(i)
+            global_score = constraint(self.observer, targets, times)
+            if constraint.boolean_constraint:
+                # This should apply to every block (if the global fails, then
+                # the local should have failed anyway)
+                constraint_zeros &= global_score
+            elif constraint.weight:
+                weight = constraint.weight
+                for i, score in enumerate(global_score):
+                    if i not in skip_global:
+                        weights[i] += weight
+                        score_array[i] += score*weight
+                        constraint_zeros[i][(score == 0)] = 0
+            else:
+                for i, score in enumerate(global_score):
+                    if i not in skip_global:
+                        weights[i] += 1
+                        score_array[i] += score
+                        constraint_zeros[i][(score == 0)] = 0
+
+        for i, scores in enumerate(score_array):
+            if weights[i]:
+                scores *= 1/float(weights[i])
+            else:
+                # if no weight, then nothing was added to the score_array
+                # just use the zeros (squaring 0 and 1 gives 0 and 1)
+                scores += constraint_zeros[i]
+        # considering the else above, score_array would be constraint_zeros^2
+        # which is fine since all of its values should be 0 or 1
+        score_array *= constraint_zeros
+        return score_array
+
     @classmethod
     def from_start_end(cls, blocks, observer, start_time, end_time,
                        global_constraints=[]):

astroplan/tests/test_scheduling.py

@@ -265,3 +265,30 @@ def test_scorer():
     scores = scorer.create_score_array(time_resolution=20 * u.minute)
     # the ``global_constraint``: constraint2 should have applied to the blocks
     assert np.array_equal(c2, scores)
+
+
+def test_weighted_scorer():
+    times = time_grid_from_range(Time(['2016-02-06 00:00', '2016-02-06 08:00']),
+                                 time_resolution=20 * u.minute)
+    constraint = AirmassConstraint(max=2, boolean_constraint=False)
+    constraint.weight = .8
+    c = constraint(apo, [vega, rigel], times)
+    block = ObservingBlock(vega, 1 * u.hour, 0, constraints=[constraint])
+    block2 = ObservingBlock(rigel, 1 * u.hour, 0, constraints=[constraint])
+    scorer = Scorer.from_start_end([block, block2], apo, Time('2016-02-06 00:00'),
+                                   Time('2016-02-06 08:00'))
+    scores = scorer.weighted_score_array(time_resolution=20 * u.minute)
+    # due to float multiplication and division c and scores are not exactly equal
+    assert np.array_equal(np.round(c, 10), np.round(scores, 10))
+
+    constraint2 = MoonIlluminationConstraint(max=.6, boolean_constraint=False)
+    constraint2.weight = .7
+    c2 = constraint2(apo, [vega, rigel], times)
+    block = ObservingBlock(vega, 1 * u.hour, 0, constraints=[constraint, constraint2])
+    block2 = ObservingBlock(rigel, 1 * u.hour, 0, constraints=[constraint])
+    scorer = Scorer.from_start_end([block, block2], apo, Time('2016-02-06 00:00'),
+                                   Time('2016-02-06 08:00'))
+    scores = scorer.weighted_score_array(time_resolution=20 * u.minute)
+    assert all(scores[0] - (c[0] * .8 + c2[0] * .7)/1.5)
+    np.array_equal(np.round(scores[0], 10), np.round((c[0] * .8 + c2[0] * .7)/1.5, 10))
+    assert np.array_equal(np.round(scores[1], 10), np.round(c[1], 10))