extend_scheduled_observations method added to Schedule (#69)

nenupy/__init__.py

@@ -5,7 +5,7 @@
 __copyright__ = "Copyright 2023, nenupy"
 __credits__ = ["Alan Loh"]
 __license__ = "MIT"
-__version__ = "2.6.11"
+__version__ = "2.7.0"
 __maintainer__ = "Alan Loh"
 __email__ = "[email protected]"
 

nenupy/schedule/constraints.py

@@ -174,6 +174,7 @@
 import pytz
 import matplotlib.pyplot as plt
 from copy import copy
+from typing import List
 
 from nenupy.schedule.targets import _Target
 
@@ -1145,6 +1146,52 @@ def evaluate(self, target, time, nslots=1, sun_elevation=None):
         )
         self.score[np.where(np.prod(cnts, axis=0)==0)[0]] = 0
         return self.score
+
+
+    def evaluate_various_nslots(self, target: _Target, time: Time, test_indices: List[np.ndarray], sun_elevation=None) -> List[np.ndarray]:
+        scores = []
+
+        # Loop over each set of indices
+        for indices_window in test_indices:
+
+            # Copy target and only keep the time/coordinates around the boundaries of test_indices to optimize the computations
+            min_idx, max_idx = np.min(indices_window), np.max(indices_window) + 1
+            sliced_target = None if target is None else target[min_idx : max_idx] # Perform a copy without damaging the original object
+            time = time[min_idx : max_idx]
+            sun_elevation = None if sun_elevation is None else sun_elevation[min_idx : max_idx]
+
+            nslots = indices_window.size - 1
+
+            # Evaluate the constraints
+            n_unevaluated_constraints = 0
+            constraint_scores = np.zeros((self.size, nslots))
+            for constraint_i, constraint in enumerate(self):
+                constraint = copy(constraint) # to prevent overwriting initial constraints
+                if isinstance(constraint, TargetConstraint) and (sliced_target is not None):
+                    constraint_scores[constraint_i, :] = constraint(sliced_target, nslots)
+                elif isinstance(constraint, ScheduleConstraint):
+                    constraint_scores[constraint_i, :] = constraint(tuple(time), nslots)
+                elif isinstance(constraint, NightTimeCnst):
+                    constraint_scores[constraint_i, :] = constraint(sun_elevation)
+                else:
+                    n_unevaluated_constraints += 1
+
+            if n_unevaluated_constraints == self.size:
+                constraint_scores += 1.
+                log.debug(
+                    "No defined constraint could be used. Schedule "
+                    "slot scores have been set to 1..."
+                )
+            # Compute the weighted average score
+            window_score = np.average(constraint_scores, weights=self.weights, axis=0)
+            # Replace NaN values by 0
+            window_score = np.where(np.isnan(window_score), 0, window_score)
+            # If any constraint is at 0, set the time-score at 0 no matter the other constraints
+            window_score[np.where(np.prod(constraint_scores, axis=0)==0)[0]] = 0
+            scores.append(window_score)
+
+        # return the scores
+        return scores
 
 
     def plot(self, **kwargs):

nenupy/schedule/obsblocks.py

@@ -339,12 +339,17 @@ class ObsBlock(Block):
             This parameter particularly suits the imaging data.
         :type processing_delay:
             :class:`~astropy.time.TimeDelta`
+        :param max_extended_duration:
+            Once the observation block is booked. The duration of the scheduled observation will be extended up to this value.
+            The duration extension will cease if the resulting scheduled score ever decreases or if another scheduled observation is reached.
+        :type max_extended_duration:
+            :class:`~astropy.time.TimeDelta`
 
         .. seealso::
             :ref:`observation_request_sec`
-        
+
         .. rubric:: Attributes Summary
-        
+
         .. autosummary::
 
             ~ObsBlock.program
@@ -359,14 +364,16 @@ def __init__(
         self, name, program, target,
         constraints=None,
         duration=TimeDelta(3600, format='sec'),
-        processing_delay: TimeDelta = None
+        processing_delay: TimeDelta = None,
+        max_extended_duration: TimeDelta = None
     ):
         self.name = name
         self.program = program
         self.target = target
         self.constraints = constraints
         self.duration = duration
         self.processing_delay = processing_delay
+        self.max_extended_duration = max_extended_duration
 
         self.blockIdx = 0
 

nenupy/schedule/schedule.py

@@ -75,15 +75,17 @@ def __init__(self,
         constraints,
         duration,
         dt,
-        processing_delay
+        processing_delay,
+        max_extended_duration
     ):
         super().__init__(
             name=name,
             program=program,
             target=target,
             constraints=constraints,
             duration=duration,
-            processing_delay=processing_delay
+            processing_delay=processing_delay,
+            max_extended_duration=max_extended_duration
         )
 
         self.dt = dt
@@ -409,7 +411,8 @@ def insert(self, block):
                     constraints=blk.constraints,
                     duration=blk.duration,
                     dt=self.dt,
-                    processing_delay=blk.processing_delay
+                    processing_delay=blk.processing_delay,
+                    max_extended_duration=blk.max_extended_duration
                 )
             sb.blockIdx = self._idxCounter
             self.blocks.append(sb)
@@ -687,6 +690,8 @@ def __init__(
         elevation = sun.elevation.deg
         self.sun_elevation = (elevation[1:] + elevation[:-1])/2
 
+        self._obs_extension_done = False
+
 
     def __getitem__(self, n):
         """
@@ -1109,6 +1114,111 @@ def book(self, optimize=False, **kwargs):
             )
 
 
+    def extend_scheduled_observations(self, relative_score_threshold: float = 0.8) -> None:
+        """ """
+        if self._obs_extension_done:
+            log.warning("Scheduled observations have already been extended.")
+            return
+
+        log.info("Extending scheduled observations...")
+
+        times = Time(
+            np.append(
+                self._startsJD,
+                self._startsJD[-1] + self.dt.jd
+            ),
+            format='jd'
+        )
+
+        # Prepare the available extensions for each observation block
+        extend_mask = []
+        slot_extension = []
+        slot_extension_indices = []
+        indices = []
+        for block in self.observation_blocks:
+            if not block.isBooked:
+                continue
+            extdt = block.max_extended_duration
+            extend_mask.append(extdt is not None)
+            n_slots = int(np.ceil(extdt.sec/self.dt.sec))
+            slot_extension.append(n_slots if extdt is not None else 0)
+            indices.append(block.blockIdx)
+
+            # Compute the constraint score for each possibility (i.e. from self.nSlots to n_slots)
+            extended_windows_size = np.arange(block.nSlots + 1, n_slots + 1) + 1 # Add 1 to evaluate the constraints
+            extended_windows_slots = [np.arange(window_size) - int(window_size/2) + block.startIdx + int(block.nSlots/2) for window_size in extended_windows_size]
+            # Deal with edges (don't check indices outside of the Schedule range)
+            #extended_windows_slots = [window_indices[(window_indices >= 0) * (window_indices < self.size - 1)] for window_indices in extended_windows_slots] # self.size - 1 to compensate for the constraints evaluation
+            extended_windows_slots = [window_indices[(window_indices >= 0) * (window_indices < self.size)] for window_indices in extended_windows_slots]
+            # Compute the constraint score
+            constraint_score = block.constraints.evaluate_various_nslots(
+                target=block.target,
+                time=times,
+                test_indices=extended_windows_slots,
+                sun_elevation=self.sun_elevation
+            )
+
+            # Compute the mean score
+            def treat(data):
+                if np.prod(data) == 0:
+                    return 0.
+                else:
+                    return np.mean(data)
+            constraint_score = np.array(list(map(treat, constraint_score)))
+            # constraint_score = np.array(list(map(np.mean, constraint_score)))
+
+            # Filter out the constraint score to discard those that diminish the block score
+            # Add a relative threshold in order to filter...
+            positive_gradient_mask = constraint_score >= constraint_score[0] * relative_score_threshold
+            slot_extension_indices.append([slot_indices for slot_indices, better_option in zip(extended_windows_slots, positive_gradient_mask) if better_option])
+
+        # Iteratively try to increase the observation blocks
+        iteration_i = 0
+        max_iteration = max([len(indices) for indices in slot_extension_indices])
+        while iteration_i < max_iteration:
+
+            booked_block_i = 0
+
+            for block in self.observation_blocks:
+
+                # If the block has not been booked, don't consider it
+                if (not block.isBooked):
+                    continue
+
+                # If the block can be extended, try to do it
+                if extend_mask[booked_block_i]:
+                    available_indices = slot_extension_indices[booked_block_i]
+                    # Iteratively try to grow
+                    try:
+                        new_block_indices = available_indices[iteration_i]
+
+                        # If the number of free slots potentially occupied by the new indices is bigger, then extend the block
+                        # + block.nSlots because the block is already scheduled and counts as occupying freeSlots
+                        if sum(self.freeSlots[new_block_indices]) + block.nSlots > block.nSlots:
+                            new_block_indices = np.array([idx for idx in new_block_indices if (self.freeSlots[idx] or (idx in block.indices))])
+
+                            # Check that the new_slots are continuous (they could be overlapping a small nearby observation)
+                            if np.all(np.diff(new_block_indices) == 1):
+                                block.startIdx = new_block_indices[0]
+                                new_block_nslots = len(new_block_indices)
+                                block.duration = new_block_nslots * self.dt
+                                block.nSlots = new_block_nslots
+                                block.time_min = self.starts[block.startIdx]
+                                block.time_max = self.stops[block.startIdx + block.nSlots - 1]
+                                self.freeSlots[block.startIdx:block.startIdx + block.nSlots] = False
+                            else:
+                                pass
+
+                    except IndexError:
+                        # Max growth reached
+                        pass
+
+                booked_block_i += 1
+
+            iteration_i += 1
+
+        self._obs_extension_done = True
+
     def fine_tune(self, max_it: int = 1000) -> None:
         """ 
         """

nenupy/schedule/targets.py

@@ -35,6 +35,7 @@
     get_body
 )
 import numpy as np
+import copy
 
 import logging
 log = logging.getLogger(__name__)
@@ -116,6 +117,16 @@ def azimuth(self):
         return self._azimuth
 
 
+    def __getitem__(self, slice_obj):
+        target_sliced = copy.copy(self)
+        target_sliced.reset()
+        target_sliced._lst = self._lst[slice_obj]
+        target_sliced._fk5 = self._fk5[slice_obj]
+        target_sliced._hourAngle = self._hourAngle[slice_obj]
+        target_sliced._elevation = self._elevation[slice_obj]
+        target_sliced._azimuth = self._azimuth[slice_obj]
+        return target_sliced
+
     # --------------------------------------------------------- #
     # ------------------------ Methods ------------------------ #
     def reset(self) -> None:
@@ -226,7 +237,7 @@ def _computeHorizontalCoords(self):
     # --------------------------------------------------------- #
     # ----------------------- Internal ------------------------ #
     @staticmethod
-    def _attrWarning(self, attr):
+    def _attrWarning(attr):
         """
         """
         log.warning(