• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

dsavransky / EXOSIMS / 13549049431

26 Feb 2025 04:48PM UTC coverage: 65.679% (-0.04%) from 65.72%
13549049431

push

github

dsavransky
trying to trick rtd into working again

9549 of 14539 relevant lines covered (65.68%)

0.66 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

61.99
/EXOSIMS/Prototypes/SurveySimulation.py
1
# -*- coding: utf-8 -*-
2
import copy
1✔
3
import hashlib
1✔
4
import inspect
1✔
5
import json
1✔
6
import logging
1✔
7
import os
1✔
8
import random as py_random
1✔
9
import re
1✔
10
import subprocess
1✔
11
import sys
1✔
12
import time
1✔
13
from pathlib import Path
1✔
14
from typing import Any, Dict, Optional
1✔
15

16
import astropy.constants as const
1✔
17
import astropy.units as u
1✔
18
import numpy as np
1✔
19
from astropy.time import Time
1✔
20

21
from EXOSIMS.util.deltaMag import deltaMag
1✔
22
from EXOSIMS.util.get_dirs import get_cache_dir
1✔
23
from EXOSIMS.util.get_module import get_module
1✔
24
from EXOSIMS.util.vprint import vprint
1✔
25
from EXOSIMS.util._numpy_compat import copy_if_needed
1✔
26

27

28
Logger = logging.getLogger(__name__)
1✔
29

30

31
class SurveySimulation(object):
1✔
32
    """:ref:`SurveySimulation` Prototype
33

34
    Args:
35

36
        scriptfile (str, optional):
37
            JSON script file.  If not set, assumes that dictionary has been
38
            passed through specs. Defaults to None\
39
        ntFlux (int):
40
            Number of intervals to split integration into for computing
41
            total SNR. When greater than 1, SNR is effectively computed as a
42
            Reimann sum. Defaults to 1
43
        nVisitsMax (int):
44
            Maximum number of observations (in detection mode) per star.
45
            Defaults to 5
46
        charMargin (float):
47
            Integration time margin for characterization. Defaults to 0.15
48
        dt_max (float):
49
            Maximum time for revisit window (in days). Defaults to 1.
50
        record_counts_path (str, optional):
51
            If set, write out photon count info to file specified by this keyword.
52
            Defaults to None.
53
        revisit_wait (float, optional):
54
            If set, it is the minimum time (in days) to wait before revisiting current target. Defaults to None.
55
        nokoMap (bool):
56
            Skip generating keepout map. Only useful if you're not planning on
57
            actually running a mission simulation. Defaults to False.
58
        nofZ (bool):
59
            Skip precomputing zodical light minima.  Only useful if you're not
60
            planning on actually running a mission simulation.  Defaults to False.
61
        cachedir (str, optional):
62
            Full path to cachedir.
63
            If None (default) use default (see :ref:`EXOSIMSCACHE`)
64
        defaultAddExoplanetObsTime (bool):
65
            If True, time advancement when no targets are observable will add
66
            to exoplanetObsTime (i.e., wasting time is counted against you).
67
            Defaults to True
68
        find_known_RV (bool):
69
            Identify stars with known planets. Defaults to False
70
        include_known_RV (str, optional):
71
            Path to file including known planets to include. Defaults to None
72
        make_debug_bird_plots (bool, optional):
73
            If True, makes completeness bird plots for every observation that
74
            are saved in the cache directory
75
        debug_plot_path (str, optional):
76
            Path to save the debug plots in, must be set if
77
            make_debug_bird_plots is True
78
        **specs:
79
            :ref:`sec:inputspec`
80

81
    Attributes:
82
        _outspec (dict):
83
            :ref:`sec:outspec`
84
        absTimefZmin (astropy.time.core.Time):
85
            Absolute time of local zodi minima
86
        BackgroundSources (:ref:`BackgroundSources`):
87
            BackgroundSources object
88
        cachedir (str):
89
            Path to the EXOSIMS cache directory (see :ref:`EXOSIMSCACHE`)
90
        cachefname (str):
91
            Base filename for cache files.
92
        charMargin (float):
93
            Integration time margin for characterization.
94
        Completeness (:ref:`Completeness`):
95
            Completeness object
96
        count_lines (list):
97
            Photon counts.  Only used when ``record_counts_path`` is set
98
        defaultAddExoplanetObsTime (bool):
99
            If True, time advancement when no targets are observable will add
100
            to exoplanetObsTime (i.e., wasting time is counted against you).
101
        DRM (list):
102
            The mission simulation.  List of observation dictionaries.
103
        dt_max (astropy.units.quantity.Quantity):
104
            Maximum time for revisit window.
105
        revisit_wait (astropy.units.quantity.Quantity):
106
            Minimum time (in days) to wait before revisiting.
107
        find_known_RV (bool):
108
            Identify stars with known planets.
109
        fullSpectra (numpy.ndarray(bool)):
110
            Array of booleans indicating whether a planet's spectrum has been
111
            fully observed.
112
        fZmins (dict):
113
            Dictionary of local zodi minimum value candidates for each observing mode
114
        fZtypes (dict):
115
            Dictionary of type of local zodi minimum candidates for each observing mode
116
        include_known_RV (str, optional):
117
            Path to file including known planets to include.
118
        intTimeFilterInds (numpy.ndarray(ind)):
119
            Indices of targets where integration times fall below cutoff value
120
        intTimesIntTimeFilter (astropy.units.quantity.Quantity):
121
            Default integration times for pre-filtering targets.
122
        known_earths (numpy.ndarray):
123
            Indices of Earth-like planets
124
        known_rocky (list):
125
            Indices of rocky planets
126
        known_stars (list):
127
            Stars with known planets
128
        koMaps (dict):
129
            Keepout Maps
130
        koTimes (astropy.time.core.Time):
131
            Times corresponding to keepout map array entries.
132
        lastDetected (numpy.ndarray):
133
            ntarg x 4. For each target, contains 4 lists with planets' detected
134
            status (boolean),
135
            exozodi brightness (in units of 1/arcsec2), delta magnitude,
136
            and working angles (in units of arcsec)
137
        lastObsTimes (astropy.units.quantity.Quantity):
138
            Contains the last observation start time for future completeness update
139
            in units of day
140
        logger (logging.Logger):
141
            Logger object
142
        modules (dict):
143
            Modules dictionary.
144
        ntFlux (int):
145
            Number of intervals to split integration into for computing
146
            total SNR. When greater than 1, SNR is effectively computed as a
147
            Reimann sum.
148
        nVisitsMax (int):
149
            Maximum number of observations (in detection mode) per star.
150
        Observatory (:ref:`Observatory`):
151
            Observatory object.
152
        OpticalSystem (:ref:`OpticalSystem`):
153
            Optical system object.
154
        partialSpectra (numpy.ndarray):
155
            Array of booleans indicating whether a planet's spectrum has been
156
            partially observed.
157
        PlanetPhysicalModel (:ref:`PlanetPhysicalModel`):
158
            Planet pysical model object
159
        PlanetPopulation (:ref:`PlanetPopulation`):
160
            Planet population object
161
        PostProcessing (:ref:`PostProcessing`):
162
            Postprocessing object
163
        propagTimes (astropy.units.quantity.Quantity):
164
            Contains the current time at each target system.
165
        record_counts_path (str, optional):
166
            If set, write out photon count info to file specified by this keyword.
167
        seed (int):
168
            Seed for random number generation
169
        SimulatedUniverse (:ref:`SimulatedUniverse`):
170
            Simulated universe object
171
        StarCatalog (:ref:`StarCatalog`):
172
            Star catalog object (only if ``keepStarCatalog`` input is True.
173
        starExtended (numpy.ndarray):
174
            TBD
175
        starRevisit (numpy.ndarray):
176
            ntargs x 2. Contains indices of targets to revisit and revisit times
177
            of these targets in units of day
178
        starVisits (numpy.ndarray):
179
            ntargs x 1. Contains the number of times each target was visited
180
        TargetList (:ref:`TargetList`):
181
            TargetList object.
182
        TimeKeeping (:ref:`TimeKeeping`):
183
            Timekeeping object
184
        valfZmin (astropy.units.quantity.Quantity):
185
            Minimum local zodi for each target.
186
        ZodiacalLight (:ref:`ZodiacalLight`):
187
            Zodiacal light object.
188

189
    """
190

191
    _modtype = "SurveySimulation"
1✔
192

193
    def __init__(
1✔
194
        self,
195
        scriptfile=None,
196
        ntFlux=1,
197
        nVisitsMax=5,
198
        charMargin=0.15,
199
        dt_max=1.0,
200
        record_counts_path=None,
201
        revisit_wait=None,
202
        nokoMap=False,
203
        nofZ=False,
204
        cachedir=None,
205
        defaultAddExoplanetObsTime=True,
206
        find_known_RV=False,
207
        include_known_RV=None,
208
        make_debug_bird_plots=False,
209
        debug_plot_path=None,
210
        **specs,
211
    ):
212

213
        # start the outspec
214
        self._outspec = {}
1✔
215

216
        # if a script file is provided read it in. If not set, assumes that
217
        # dictionary has been passed through specs.
218
        if scriptfile is not None:
1✔
219
            assert os.path.isfile(scriptfile), "%s is not a file." % scriptfile
1✔
220

221
            try:
1✔
222
                with open(scriptfile) as ff:
1✔
223
                    script = ff.read()
1✔
224
                specs.update(json.loads(script))
1✔
225
            except ValueError:
1✔
226
                sys.stderr.write("Script file `%s' is not valid JSON." % scriptfile)
1✔
227
                # must re-raise, or the error will be masked
228
                raise
1✔
229

230
            # modules array must be present
231
            if "modules" not in specs:
1✔
232
                raise ValueError("No modules field found in script.")
×
233

234
        # load the vprint function (same line in all prototype module constructors)
235
        self.vprint = vprint(specs.get("verbose", True))
1✔
236

237
        # count dict contains all of the C info for each star index
238
        self.record_counts_path = record_counts_path
1✔
239
        self.count_lines = []
1✔
240
        self._outspec["record_counts_path"] = record_counts_path
1✔
241

242
        # mission simulation logger
243
        self.logger = specs.get("logger", logging.getLogger(__name__))
1✔
244

245
        # set up numpy random number (generate it if not in specs)
246
        self.seed = int(specs.get("seed", py_random.randint(1, int(1e9))))
1✔
247
        self.vprint("Numpy random seed is: %s" % self.seed)
1✔
248
        np.random.seed(self.seed)
1✔
249
        self._outspec["seed"] = self.seed
1✔
250

251
        # cache directory
252
        self.cachedir = get_cache_dir(cachedir)
1✔
253
        self._outspec["cachedir"] = self.cachedir
1✔
254
        # N.B.: cachedir is going to be used by everything, so let's make sure that
255
        # it doesn't get popped out of specs
256
        specs["cachedir"] = self.cachedir
1✔
257

258
        # if any of the modules is a string, assume that they are all strings
259
        # and we need to initalize
260
        if isinstance(next(iter(specs["modules"].values())), str):
1✔
261

262
            # import desired module names (prototype or specific)
263
            self.SimulatedUniverse = get_module(
1✔
264
                specs["modules"]["SimulatedUniverse"], "SimulatedUniverse"
265
            )(**specs)
266
            self.Observatory = get_module(
1✔
267
                specs["modules"]["Observatory"], "Observatory"
268
            )(**specs)
269
            self.TimeKeeping = get_module(
1✔
270
                specs["modules"]["TimeKeeping"], "TimeKeeping"
271
            )(**specs)
272

273
            # bring inherited class objects to top level of Survey Simulation
274
            SU = self.SimulatedUniverse
1✔
275
            self.StarCatalog = SU.StarCatalog
1✔
276
            self.PlanetPopulation = SU.PlanetPopulation
1✔
277
            self.PlanetPhysicalModel = SU.PlanetPhysicalModel
1✔
278
            self.OpticalSystem = SU.OpticalSystem
1✔
279
            self.ZodiacalLight = SU.ZodiacalLight
1✔
280
            self.BackgroundSources = SU.BackgroundSources
1✔
281
            self.PostProcessing = SU.PostProcessing
1✔
282
            self.Completeness = SU.Completeness
1✔
283
            self.TargetList = SU.TargetList
1✔
284

285
        else:
286
            # these are the modules that must be present if passing instantiated objects
287
            neededObjMods = [
1✔
288
                "PlanetPopulation",
289
                "PlanetPhysicalModel",
290
                "OpticalSystem",
291
                "ZodiacalLight",
292
                "BackgroundSources",
293
                "PostProcessing",
294
                "Completeness",
295
                "TargetList",
296
                "SimulatedUniverse",
297
                "Observatory",
298
                "TimeKeeping",
299
            ]
300

301
            # ensure that you have the minimal set
302
            for modName in neededObjMods:
1✔
303
                if modName not in specs["modules"]:
1✔
304
                    raise ValueError(
×
305
                        "%s module is required but was not provided." % modName
306
                    )
307

308
            for modName in specs["modules"]:
1✔
309
                assert specs["modules"][modName]._modtype == modName, (
1✔
310
                    "Provided instance of %s has incorrect modtype." % modName
311
                )
312

313
                setattr(self, modName, specs["modules"][modName])
1✔
314

315
        # create a dictionary of all modules, except StarCatalog
316
        self.modules = {}
1✔
317
        self.modules["PlanetPopulation"] = self.PlanetPopulation
1✔
318
        self.modules["PlanetPhysicalModel"] = self.PlanetPhysicalModel
1✔
319
        self.modules["OpticalSystem"] = self.OpticalSystem
1✔
320
        self.modules["ZodiacalLight"] = self.ZodiacalLight
1✔
321
        self.modules["BackgroundSources"] = self.BackgroundSources
1✔
322
        self.modules["PostProcessing"] = self.PostProcessing
1✔
323
        self.modules["Completeness"] = self.Completeness
1✔
324
        self.modules["TargetList"] = self.TargetList
1✔
325
        self.modules["SimulatedUniverse"] = self.SimulatedUniverse
1✔
326
        self.modules["Observatory"] = self.Observatory
1✔
327
        self.modules["TimeKeeping"] = self.TimeKeeping
1✔
328
        # add yourself to modules list for bookkeeping purposes
329
        self.modules["SurveySimulation"] = self
1✔
330

331
        # observation time sampling
332
        self.ntFlux = int(ntFlux)
1✔
333
        self._outspec["ntFlux"] = self.ntFlux
1✔
334

335
        # maximum number of observations per star
336
        self.nVisitsMax = int(nVisitsMax)
1✔
337
        self._outspec["nVisitsMax"] = self.nVisitsMax
1✔
338

339
        # integration time margin for characterization
340
        self.charMargin = float(charMargin)
1✔
341
        self._outspec["charMargin"] = self.charMargin
1✔
342

343
        # maximum time for revisit window
344
        self.dt_max = float(dt_max) * u.week
1✔
345
        self._outspec["dt_max"] = self.dt_max.value
1✔
346

347
        # minimum time for revisit window
348
        if revisit_wait is not None:
1✔
349
            self.revisit_wait = revisit_wait * u.d
×
350
        else:
351
            self.revisit_wait = revisit_wait
1✔
352

353
        # list of detected earth-like planets aroung promoted stars
354
        self.known_earths = np.array([])
1✔
355

356
        self.find_known_RV = find_known_RV
1✔
357
        self._outspec["find_known_RV"] = find_known_RV
1✔
358
        self._outspec["include_known_RV"] = include_known_RV
1✔
359
        if self.find_known_RV:
1✔
360
            # select specific knonw RV stars if a file exists
361
            if include_known_RV is not None:
×
362
                if os.path.isfile(include_known_RV):
×
363
                    with open(include_known_RV, "r") as rv_file:
×
364
                        self.include_known_RV = [
×
365
                            hip.strip() for hip in rv_file.read().split(",")
366
                        ]
367
                        self.vprint(
×
368
                            "Including known RV stars: {}".format(self.include_known_RV)
369
                        )
370
                else:
371
                    self.include_known_RV = None
×
372
                    self.vprint(
×
373
                        "WARNING: Known RV file: {} does not exist!!".format(
374
                            include_known_RV
375
                        )
376
                    )
377
            else:
378
                self.include_known_RV = None
×
379
            self.known_stars, self.known_rocky = self.find_known_plans()
×
380
        else:
381
            self.include_known_RV = None
1✔
382
            self.known_stars = []
1✔
383
            self.known_rocky = []
1✔
384

385
        # defaultAddExoplanetObsTime Tells us time advanced when no targets available
386
        # counts agains exoplanetObsTime (when True)
387
        self.defaultAddExoplanetObsTime = defaultAddExoplanetObsTime
1✔
388
        self._outspec["defaultAddExoplanetObsTime"] = defaultAddExoplanetObsTime
1✔
389

390
        # If inputs are scalars, save scalars to outspec, otherwise save full lists
391
        OS = self.OpticalSystem
1✔
392
        TL = self.TargetList
1✔
393
        SU = self.SimulatedUniverse
1✔
394
        TK = self.TimeKeeping
1✔
395

396
        # initialize arrays updated in run_sim()
397
        self.initializeStorageArrays()
1✔
398

399
        # Generate File Hashnames and loction
400
        self.cachefname = self.generateHashfName(specs)
1✔
401

402
        # choose observing modes selected for detection (default marked with a flag)
403
        allModes = OS.observingModes
1✔
404
        det_mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
405

406
        # getting keepout map for entire mission
407
        startTime = self.TimeKeeping.missionStart.copy()
1✔
408
        endTime = self.TimeKeeping.missionFinishAbs.copy()
1✔
409

410
        nSystems = len(allModes)
1✔
411
        systNames = np.unique(
1✔
412
            [allModes[x]["syst"]["name"] for x in np.arange(nSystems)]
413
        )
414
        systOrder = np.argsort(systNames)
1✔
415
        koStr = ["koAngles_Sun", "koAngles_Moon", "koAngles_Earth", "koAngles_Small"]
1✔
416
        koangles = np.zeros([len(systNames), 4, 2])
1✔
417

418
        for x in systOrder:
1✔
419
            rel_mode = list(
1✔
420
                filter(lambda mode: mode["syst"]["name"] == systNames[x], allModes)
421
            )[0]
422
            koangles[x] = np.asarray([rel_mode["syst"][k] for k in koStr])
1✔
423

424
        self._outspec["nokoMap"] = nokoMap
1✔
425
        if not (nokoMap):
1✔
426
            koMaps, self.koTimes = self.Observatory.generate_koMap(
1✔
427
                TL, startTime, endTime, koangles
428
            )
429
            self.koMaps = {}
1✔
430
            for x, n in zip(systOrder, systNames[systOrder]):
1✔
431
                print(n)
1✔
432
                self.koMaps[n] = koMaps[x, :, :]
1✔
433

434
        self._outspec["nofZ"] = nofZ
1✔
435

436
        # TODO: do we still want a nofZ option?  probably not.
437
        self.fZmins = {}
1✔
438
        self.fZtypes = {}
1✔
439
        for x, n in zip(systOrder, systNames[systOrder]):
1✔
440
            self.fZmins[n] = np.array([])
1✔
441
            self.fZtypes[n] = np.array([])
1✔
442

443
        # TODO: do we need to do this for all modes? doing det only breaks other
444
        # schedulers, but maybe there's a better approach here.
445
        sInds = np.arange(TL.nStars)  # Initialize some sInds array
1✔
446
        for mode in allModes:
1✔
447
            # This instantiates fZMap arrays for every starlight suppresion system
448
            # that is actually used in a mode
449
            modeHashName = (
1✔
450
                f'{self.cachefname[0:-1]}_{TL.nStars}_{mode["syst"]["name"]}.'
451
            )
452

453
            if (np.size(self.fZmins[mode["syst"]["name"]]) == 0) or (
1✔
454
                np.size(self.fZtypes[mode["syst"]["name"]]) == 0
455
            ):
456
                self.ZodiacalLight.generate_fZ(
1✔
457
                    self.Observatory,
458
                    TL,
459
                    self.TimeKeeping,
460
                    mode,
461
                    modeHashName,
462
                    self.koTimes,
463
                )
464

465
                (
1✔
466
                    self.fZmins[mode["syst"]["name"]],
467
                    self.fZtypes[mode["syst"]["name"]],
468
                ) = self.ZodiacalLight.calcfZmin(
469
                    sInds,
470
                    self.Observatory,
471
                    TL,
472
                    self.TimeKeeping,
473
                    mode,
474
                    modeHashName,
475
                    self.koMaps[mode["syst"]["name"]],
476
                    self.koTimes,
477
                )
478

479
        # Precalculating intTimeFilter for coronagraph
480
        # find fZmin to use in intTimeFilter
481
        self.valfZmin, self.absTimefZmin = self.ZodiacalLight.extractfZmin(
1✔
482
            self.fZmins[det_mode["syst"]["name"]], sInds, self.koTimes
483
        )
484
        fEZ = self.ZodiacalLight.fEZ0  # grabbing fEZ0
1✔
485
        dMag = TL.int_dMag[sInds]  # grabbing dMag
1✔
486
        WA = TL.int_WA[sInds]  # grabbing WA
1✔
487
        self.intTimesIntTimeFilter = (
1✔
488
            self.OpticalSystem.calc_intTime(
489
                TL, sInds, self.valfZmin, fEZ, dMag, WA, det_mode, TK=TK
490
            )
491
            * det_mode["timeMultiplier"]
492
        )  # intTimes to filter by
493
        # These indices are acceptable for use simulating
494
        self.intTimeFilterInds = np.where(
1✔
495
            (
496
                (self.intTimesIntTimeFilter > 0)
497
                & (self.intTimesIntTimeFilter <= self.OpticalSystem.intCutoff)
498
            )
499
        )[0]
500

501
        self.make_debug_bird_plots = make_debug_bird_plots
1✔
502
        if self.make_debug_bird_plots:
1✔
503

504
            assert (
×
505
                debug_plot_path is not None
506
            ), "debug_plot_path must be set by input if make_debug_bird_plots is True"
507
            self.obs_plot_path = Path(f"{debug_plot_path}/{self.seed}")
×
508
            # Make directory if it doesn't exist
509
            if not self.obs_plot_path.exists():
×
510
                vprint(f"Making plot directory: {self.obs_plot_path}")
×
511
                self.obs_plot_path.mkdir(parents=True, exist_ok=True)
×
512
            self.obs_n_counter = 0
×
513

514
    def initializeStorageArrays(self):
1✔
515
        """
516
        Initialize all storage arrays based on # of stars and targets
517
        """
518

519
        self.DRM = []
1✔
520
        self.fullSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
1✔
521
        self.partialSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
1✔
522
        self.propagTimes = np.zeros(self.TargetList.nStars) * u.d
1✔
523
        self.lastObsTimes = np.zeros(self.TargetList.nStars) * u.d
1✔
524
        self.starVisits = np.zeros(
1✔
525
            self.TargetList.nStars, dtype=int
526
        )  # contains the number of times each star was visited
527
        self.starRevisit = np.array([])
1✔
528
        self.starExtended = np.array([], dtype=int)
1✔
529
        self.lastDetected = np.empty((self.TargetList.nStars, 4), dtype=object)
1✔
530

531
    def __str__(self):
1✔
532
        """String representation of the Survey Simulation object
533

534
        When the command 'print' is used on the Survey Simulation object, this
535
        method will return the values contained in the object
536

537
        """
538

539
        for att in self.__dict__:
1✔
540
            print("%s: %r" % (att, getattr(self, att)))
1✔
541

542
        return "Survey Simulation class object attributes"
1✔
543

544
    def run_sim(self):
1✔
545
        """Performs the survey simulation"""
546

547
        OS = self.OpticalSystem
1✔
548
        TL = self.TargetList
1✔
549
        SU = self.SimulatedUniverse
1✔
550
        Obs = self.Observatory
1✔
551
        TK = self.TimeKeeping
1✔
552

553
        # TODO: start using this self.currentSep
554
        # set occulter separation if haveOcculter
555
        if OS.haveOcculter:
1✔
556
            self.currentSep = Obs.occulterSep
×
557

558
        # choose observing modes selected for detection (default marked with a flag)
559
        allModes = OS.observingModes
1✔
560
        det_mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
561
        # and for characterization (default is first spectro/IFS mode)
562
        spectroModes = list(
1✔
563
            filter(lambda mode: "spec" in mode["inst"]["name"], allModes)
564
        )
565
        if np.any(spectroModes):
1✔
566
            char_mode = spectroModes[0]
×
567
        # if no spectro mode, default char mode is first observing mode
568
        else:
569
            char_mode = allModes[0]
1✔
570

571
        # begin Survey, and loop until mission is finished
572
        log_begin = "OB%s: survey beginning." % (TK.OBnumber)
1✔
573
        self.logger.info(log_begin)
1✔
574
        self.vprint(log_begin)
1✔
575
        t0 = time.time()
1✔
576
        sInd = None
1✔
577
        ObsNum = 0
1✔
578
        while not TK.mission_is_over(OS, Obs, det_mode):
1✔
579

580
            # acquire the NEXT TARGET star index and create DRM
581
            old_sInd = sInd  # used to save sInd if returned sInd is None
1✔
582
            DRM, sInd, det_intTime, waitTime = self.next_target(sInd, det_mode)
1✔
583

584
            if sInd is not None:
1✔
585
                ObsNum += (
1✔
586
                    1  # we're making an observation so increment observation number
587
                )
588

589
                if OS.haveOcculter:
1✔
590
                    # advance to start of observation (add slew time for selected target
591
                    _ = TK.advanceToAbsTime(TK.currentTimeAbs.copy() + waitTime)
×
592

593
                # beginning of observation, start to populate DRM
594
                DRM["star_ind"] = sInd
1✔
595
                DRM["star_name"] = TL.Name[sInd]
1✔
596
                DRM["arrival_time"] = TK.currentTimeNorm.to("day").copy()
1✔
597
                DRM["OB_nb"] = TK.OBnumber
1✔
598
                DRM["ObsNum"] = ObsNum
1✔
599
                pInds = np.where(SU.plan2star == sInd)[0]
1✔
600
                DRM["plan_inds"] = pInds.astype(int)
1✔
601
                log_obs = (
1✔
602
                    "  Observation #%s, star ind %s (of %s) with %s planet(s), "
603
                    + "mission time at Obs start: %s, exoplanetObsTime: %s"
604
                ) % (
605
                    ObsNum,
606
                    sInd,
607
                    TL.nStars,
608
                    len(pInds),
609
                    TK.currentTimeNorm.to("day").copy().round(2),
610
                    TK.exoplanetObsTime.to("day").copy().round(2),
611
                )
612
                self.logger.info(log_obs)
1✔
613
                self.vprint(log_obs)
1✔
614

615
                # PERFORM DETECTION and populate revisit list attribute
616
                (
1✔
617
                    detected,
618
                    det_fZ,
619
                    det_systemParams,
620
                    det_SNR,
621
                    FA,
622
                ) = self.observation_detection(sInd, det_intTime.copy(), det_mode)
623
                # update the occulter wet mass
624
                if OS.haveOcculter:
1✔
625
                    DRM = self.update_occulter_mass(
×
626
                        DRM, sInd, det_intTime.copy(), "det"
627
                    )
628
                # populate the DRM with detection results
629
                DRM["det_time"] = det_intTime.to("day")
1✔
630
                DRM["det_status"] = detected
1✔
631
                DRM["det_SNR"] = det_SNR
1✔
632
                DRM["det_fZ"] = det_fZ.to("1/arcsec2")
1✔
633
                DRM["det_params"] = det_systemParams
1✔
634

635
                # PERFORM CHARACTERIZATION and populate spectra list attribute
636
                if char_mode["SNR"] not in [0, np.inf]:
1✔
637
                    (
1✔
638
                        characterized,
639
                        char_fZ,
640
                        char_systemParams,
641
                        char_SNR,
642
                        char_intTime,
643
                    ) = self.observation_characterization(sInd, char_mode)
644
                else:
645
                    char_intTime = None
×
646
                    lenChar = len(pInds) + 1 if FA else len(pInds)
×
647
                    characterized = np.zeros(lenChar, dtype=float)
×
648
                    char_SNR = np.zeros(lenChar, dtype=float)
×
649
                    char_fZ = 0.0 / u.arcsec**2
×
650
                    char_systemParams = SU.dump_system_params(sInd)
×
651
                assert char_intTime != 0, "Integration time can't be 0."
1✔
652
                # update the occulter wet mass
653
                if OS.haveOcculter and (char_intTime is not None):
1✔
654
                    DRM = self.update_occulter_mass(DRM, sInd, char_intTime, "char")
×
655
                # populate the DRM with characterization results
656
                DRM["char_time"] = (
1✔
657
                    char_intTime.to("day") if char_intTime is not None else 0.0 * u.day
658
                )
659
                DRM["char_status"] = characterized[:-1] if FA else characterized
1✔
660
                DRM["char_SNR"] = char_SNR[:-1] if FA else char_SNR
1✔
661
                DRM["char_fZ"] = char_fZ.to("1/arcsec2")
1✔
662
                DRM["char_params"] = char_systemParams
1✔
663
                # populate the DRM with FA results
664
                DRM["FA_det_status"] = int(FA)
1✔
665
                DRM["FA_char_status"] = characterized[-1] if FA else 0
1✔
666
                DRM["FA_char_SNR"] = char_SNR[-1] if FA else 0.0
1✔
667
                DRM["FA_char_fEZ"] = (
1✔
668
                    self.lastDetected[sInd, 1][-1] / u.arcsec**2
669
                    if FA
670
                    else 0.0 / u.arcsec**2
671
                )
672
                DRM["FA_char_dMag"] = self.lastDetected[sInd, 2][-1] if FA else 0.0
1✔
673
                DRM["FA_char_WA"] = (
1✔
674
                    self.lastDetected[sInd, 3][-1] * u.arcsec if FA else 0.0 * u.arcsec
675
                )
676

677
                # populate the DRM with observation modes
678
                DRM["det_mode"] = dict(det_mode)
1✔
679
                del DRM["det_mode"]["inst"], DRM["det_mode"]["syst"]
1✔
680
                DRM["char_mode"] = dict(char_mode)
1✔
681
                del DRM["char_mode"]["inst"], DRM["char_mode"]["syst"]
1✔
682
                DRM["exoplanetObsTime"] = TK.exoplanetObsTime.copy()
1✔
683

684
                # append result values to self.DRM
685
                self.DRM.append(DRM)
1✔
686

687
                # handle case of inf OBs and missionPortion < 1
688
                if np.isinf(TK.OBduration) and (TK.missionPortion < 1.0):
1✔
689
                    self.arbitrary_time_advancement(
1✔
690
                        TK.currentTimeNorm.to("day").copy() - DRM["arrival_time"]
691
                    )
692

693
            else:  # sInd == None
694
                sInd = old_sInd  # Retain the last observed star
×
695
                if (
×
696
                    TK.currentTimeNorm.copy() >= TK.OBendTimes[TK.OBnumber]
697
                ):  # currentTime is at end of OB
698
                    # Conditional Advance To Start of Next OB
699
                    if not TK.mission_is_over(
×
700
                        OS, Obs, det_mode
701
                    ):  # as long as the mission is not over
702
                        TK.advancetToStartOfNextOB()  # Advance To Start of Next OB
×
703
                elif waitTime is not None:
×
704
                    # CASE 1: Advance specific wait time
705
                    _ = TK.advanceToAbsTime(
×
706
                        TK.currentTimeAbs.copy() + waitTime,
707
                        self.defaultAddExoplanetObsTime,
708
                    )
709
                    self.vprint("waitTime is not None")
×
710
                else:
711
                    startTimes = (
×
712
                        TK.currentTimeAbs.copy() + np.zeros(TL.nStars) * u.d
713
                    )  # Start Times of Observations
714
                    observableTimes = Obs.calculate_observableTimes(
×
715
                        TL,
716
                        np.arange(TL.nStars),
717
                        startTimes,
718
                        self.koMaps,
719
                        self.koTimes,
720
                        det_mode,
721
                    )[0]
722
                    # CASE 2 If There are no observable targets for the rest of the
723
                    # mission
724
                    if (
×
725
                        observableTimes[
726
                            (
727
                                TK.missionFinishAbs.copy().value * u.d
728
                                > observableTimes.value * u.d
729
                            )
730
                            * (
731
                                observableTimes.value * u.d
732
                                >= TK.currentTimeAbs.copy().value * u.d
733
                            )
734
                        ].shape[0]
735
                    ) == 0:
736
                        self.vprint(
×
737
                            (
738
                                "No Observable Targets for Remainder of mission at "
739
                                "currentTimeNorm = {}"
740
                            ).format(TK.currentTimeNorm)
741
                        )
742
                        # Manually advancing time to mission end
743
                        TK.currentTimeNorm = TK.missionLife
×
744
                        TK.currentTimeAbs = TK.missionFinishAbs
×
745
                    # CASE 3 nominal wait time if at least 1 target is still in list
746
                    # and observable
747
                    else:
748
                        # TODO: ADD ADVANCE TO WHEN FZMIN OCURS
749
                        inds1 = np.arange(TL.nStars)[
×
750
                            observableTimes.value * u.d
751
                            > TK.currentTimeAbs.copy().value * u.d
752
                        ]
753
                        # apply intTime filter
754
                        inds2 = np.intersect1d(self.intTimeFilterInds, inds1)
×
755
                        # apply revisit Filter #NOTE this means stars you added to the
756
                        # revisit list
757
                        inds3 = self.revisitFilter(
×
758
                            inds2, TK.currentTimeNorm.copy() + self.dt_max.to(u.d)
759
                        )
760
                        self.vprint(
×
761
                            "Filtering %d stars from advanceToAbsTime"
762
                            % (TL.nStars - len(inds3))
763
                        )
764
                        oTnowToEnd = observableTimes[inds3]
×
765
                        # there is at least one observableTime between now and the end
766
                        # of the mission
767
                        if not oTnowToEnd.value.shape[0] == 0:
×
768
                            # advance to that observable time
769
                            tAbs = np.min(oTnowToEnd)
×
770
                        else:
771
                            tAbs = (
×
772
                                TK.missionStart + TK.missionLife
773
                            )  # advance to end of mission
774
                        tmpcurrentTimeNorm = TK.currentTimeNorm.copy()
×
775
                        # Advance Time to this time OR start of next
776
                        # OB following this time
777
                        _ = TK.advanceToAbsTime(tAbs, self.defaultAddExoplanetObsTime)
×
778
                        self.vprint(
×
779
                            (
780
                                "No Observable Targets a currentTimeNorm= {:.2f}. "
781
                                "Advanced To {:.2f}"
782
                            ).format(
783
                                tmpcurrentTimeNorm.to("day"),
784
                                TK.currentTimeNorm.to("day"),
785
                            )
786
                        )
787
        else:  # TK.mission_is_over()
788
            dtsim = (time.time() - t0) * u.s
1✔
789
            log_end = (
1✔
790
                "Mission complete: no more time available.\n"
791
                + "Simulation duration: %s.\n" % dtsim.astype("int")
792
                + "Results stored in SurveySimulation.DRM (Design Reference Mission)."
793
            )
794
            self.logger.info(log_end)
1✔
795
            self.vprint(log_end)
1✔
796

797
    def arbitrary_time_advancement(self, dt):
1✔
798
        """Handles fully dynamically scheduled case where OBduration is infinite and
799
        missionPortion is less than 1.
800

801
        Args:
802
            dt (~astropy.units.quantity.Quantity):
803
                Total amount of time, including all overheads and extras used for the
804
                previous observation.
805

806
        Returns:
807
            None
808
        """
809

810
        self.TimeKeeping.allocate_time(
1✔
811
            dt
812
            * (1.0 - self.TimeKeeping.missionPortion)
813
            / self.TimeKeeping.missionPortion,
814
            addExoplanetObsTime=False,
815
        )
816

817
    def next_target(self, old_sInd, mode):
1✔
818
        """Finds index of next target star and calculates its integration time.
819

820
        This method chooses the next target star index based on which
821
        stars are available, their integration time, and maximum completeness.
822
        Returns None if no target could be found.
823

824
        Args:
825
            old_sInd (int):
826
                Index of the previous target star
827
            mode (dict):
828
                Selected observing mode for detection
829

830
        Returns:
831
            tuple:
832
                DRM (dict):
833
                    Design Reference Mission, contains the results of one complete
834
                    observation (detection and characterization)
835
                sInd (int):
836
                    Index of next target star. Defaults to None.
837
                intTime (astropy.units.Quantity):
838
                    Selected star integration time for detection in units of day.
839
                    Defaults to None.
840
                waitTime (astropy.units.Quantity):
841
                    a strategically advantageous amount of time to wait in the case
842
                    of an occulter for slew times
843

844
        """
845
        OS = self.OpticalSystem
1✔
846
        TL = self.TargetList
1✔
847
        Obs = self.Observatory
1✔
848
        TK = self.TimeKeeping
1✔
849

850
        # create DRM
851
        DRM = {}
1✔
852

853
        # allocate settling time + overhead time
854
        tmpCurrentTimeAbs = (
1✔
855
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
856
        )
857
        tmpCurrentTimeNorm = (
1✔
858
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
859
        )
860

861
        # create appropriate koMap
862
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
863

864
        # look for available targets
865
        # 1. initialize arrays
866
        slewTimes = np.zeros(TL.nStars) * u.d
1✔
867
        # fZs = np.zeros(TL.nStars) / u.arcsec**2.0
868
        dV = np.zeros(TL.nStars) * u.m / u.s
1✔
869
        intTimes = np.zeros(TL.nStars) * u.d
1✔
870
        obsTimes = np.zeros([2, TL.nStars]) * u.d
1✔
871
        sInds = np.arange(TL.nStars)
1✔
872

873
        # 2. find spacecraft orbital START positions (if occulter, positions
874
        # differ for each star) and filter out unavailable targets
875
        sd = None
1✔
876
        if OS.haveOcculter:
1✔
877
            sd = Obs.star_angularSep(TL, old_sInd, sInds, tmpCurrentTimeAbs)
×
878
            obsTimes = Obs.calculate_observableTimes(
×
879
                TL, sInds, tmpCurrentTimeAbs, self.koMaps, self.koTimes, mode
880
            )
881
            slewTimes = Obs.calculate_slewTimes(
×
882
                TL, old_sInd, sInds, sd, obsTimes, tmpCurrentTimeAbs
883
            )
884

885
        # 2.1 filter out totTimes > integration cutoff
886
        if len(sInds.tolist()) > 0:
1✔
887
            sInds = np.intersect1d(self.intTimeFilterInds, sInds)
1✔
888

889
        # start times, including slew times
890
        startTimes = tmpCurrentTimeAbs.copy() + slewTimes
1✔
891
        startTimesNorm = tmpCurrentTimeNorm.copy() + slewTimes
1✔
892

893
        # 2.5 Filter stars not observable at startTimes
894
        try:
1✔
895
            tmpIndsbool = list()
1✔
896
            for i in np.arange(len(sInds)):
1✔
897
                koTimeInd = np.where(
1✔
898
                    np.round(startTimes[sInds[i]].value) - self.koTimes.value == 0
899
                )[0][
900
                    0
901
                ]  # find indice where koTime is startTime[0]
902
                tmpIndsbool.append(
1✔
903
                    koMap[sInds[i]][koTimeInd].astype(bool)
904
                )  # Is star observable at time ind
905
            sInds = sInds[tmpIndsbool]
1✔
906
            del tmpIndsbool
1✔
907
        except:  # noqa: E722 # If there are no target stars to observe
×
908
            sInds = np.asarray([], dtype=int)
×
909

910
        # 3. filter out all previously (more-)visited targets, unless in
911
        if len(sInds.tolist()) > 0:
1✔
912
            sInds = self.revisitFilter(sInds, tmpCurrentTimeNorm)
1✔
913

914
        # 4.1 calculate integration times for ALL preselected targets
915
        (
1✔
916
            maxIntTimeOBendTime,
917
            maxIntTimeExoplanetObsTime,
918
            maxIntTimeMissionLife,
919
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
920
        maxIntTime = min(
1✔
921
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
922
        )  # Maximum intTime allowed
923

924
        if len(sInds.tolist()) > 0:
1✔
925
            if OS.haveOcculter and (old_sInd is not None):
1✔
926
                (
×
927
                    sInds,
928
                    slewTimes[sInds],
929
                    intTimes[sInds],
930
                    dV[sInds],
931
                ) = self.refineOcculterSlews(
932
                    old_sInd, sInds, slewTimes, obsTimes, sd, mode
933
                )
934
                endTimes = tmpCurrentTimeAbs.copy() + intTimes + slewTimes
×
935
            else:
936
                intTimes[sInds] = self.calc_targ_intTime(sInds, startTimes[sInds], mode)
1✔
937
                sInds = sInds[
1✔
938
                    np.where(intTimes[sInds] <= maxIntTime)
939
                ]  # Filters targets exceeding end of OB
940
                endTimes = tmpCurrentTimeAbs.copy() + intTimes
1✔
941

942
                if maxIntTime.value <= 0:
1✔
943
                    sInds = np.asarray([], dtype=int)
×
944

945
        # 5.1 TODO Add filter to filter out stars entering and exiting keepout
946
        # between startTimes and endTimes
947

948
        # 5.2 find spacecraft orbital END positions (for each candidate target),
949
        # and filter out unavailable targets
950
        if len(sInds.tolist()) > 0 and Obs.checkKeepoutEnd:
1✔
951
            # endTimes may exist past koTimes so we have an exception to hand this case
952
            try:
1✔
953
                tmpIndsbool = list()
1✔
954
                for i in np.arange(len(sInds)):
1✔
955
                    koTimeInd = np.where(
1✔
956
                        np.round(endTimes[sInds[i]].value) - self.koTimes.value == 0
957
                    )[0][
958
                        0
959
                    ]  # find indice where koTime is endTime[0]
960
                    tmpIndsbool.append(
1✔
961
                        koMap[sInds[i]][koTimeInd].astype(bool)
962
                    )  # Is star observable at time ind
963
                sInds = sInds[tmpIndsbool]
1✔
964
                del tmpIndsbool
1✔
965
            except:  # noqa: E722
×
966
                sInds = np.asarray([], dtype=int)
×
967

968
        # 6. choose best target from remaining
969
        if len(sInds.tolist()) > 0:
1✔
970
            # choose sInd of next target
971
            sInd, waitTime = self.choose_next_target(
1✔
972
                old_sInd, sInds, slewTimes, intTimes[sInds]
973
            )
974

975
            # Should Choose Next Target decide there are no stars it wishes to
976
            # observe at this time.
977
            if sInd is None and (waitTime is not None):
1✔
978
                self.vprint(
×
979
                    "There are no stars available to observe. Waiting {}".format(
980
                        waitTime
981
                    )
982
                )
983
                return DRM, None, None, waitTime
×
984
            elif (sInd is None) and (waitTime is None):
1✔
985
                self.vprint(
×
986
                    "There are no stars available to observe and waitTime is None."
987
                )
988
                return DRM, None, None, waitTime
×
989
            # store selected star integration time
990
            intTime = intTimes[sInd]
1✔
991

992
        # if no observable target, advanceTime to next Observable Target
993
        else:
994
            self.vprint(
×
995
                "No Observable Targets at currentTimeNorm= "
996
                + str(TK.currentTimeNorm.copy())
997
            )
998
            return DRM, None, None, None
×
999

1000
        # update visited list for selected star
1001
        self.starVisits[sInd] += 1
1✔
1002
        # store normalized start time for future completeness update
1003
        self.lastObsTimes[sInd] = startTimesNorm[sInd]
1✔
1004

1005
        # populate DRM with occulter related values
1006
        if OS.haveOcculter:
1✔
1007
            DRM = Obs.log_occulterResults(
×
1008
                DRM, slewTimes[sInd], sInd, sd[sInd], dV[sInd]
1009
            )
1010
            return DRM, sInd, intTime, slewTimes[sInd]
×
1011

1012
        return DRM, sInd, intTime, waitTime
1✔
1013

1014
    def calc_targ_intTime(self, sInds, startTimes, mode):
1✔
1015
        """Helper method for next_target to aid in overloading for alternative
1016
        implementations.
1017

1018
        Given a subset of targets, calculate their integration times given the
1019
        start of observation time.
1020

1021
        Prototype just calculates integration times for fixed contrast depth.
1022

1023
        Args:
1024
            sInds (~numpy.ndarray(int)):
1025
                Indices of available targets
1026
            startTimes (astropy quantity array):
1027
                absolute start times of observations.
1028
                must be of the same size as sInds
1029
            mode (dict):
1030
                Selected observing mode for detection
1031

1032
        Returns:
1033
            ~astropy.units.Quantity(~numpy.ndarray(float)):
1034
                Integration times for detection
1035
                same dimension as sInds
1036

1037
        .. note::
1038

1039
            next_target filter will discard targets with zero integration times.
1040

1041
        """
1042

1043
        SU = self.SimulatedUniverse
1✔
1044
        TL = self.TargetList
1✔
1045

1046
        # assumed values for detection
1047
        fZ = self.ZodiacalLight.fZ(
1✔
1048
            self.Observatory, self.TargetList, sInds, startTimes, mode
1049
        )
1050
        fEZ = self.ZodiacalLight.fEZ0
1✔
1051
        fEZs = np.zeros(len(sInds)) / u.arcsec**2
1✔
1052
        for i, sInd in enumerate(sInds):
1✔
1053
            pInds = np.where(SU.plan2star == sInd)[0]
1✔
1054
            pInds_earthlike = pInds[self.is_earthlike(pInds, sInd)]
1✔
1055
            if len(pInds_earthlike) == 0:
1✔
1056
                fEZs[i] = fEZ
1✔
1057
            else:
1058
                fEZs[i] = np.max(SU.fEZ[pInds_earthlike])
1✔
1059
        dMag = TL.int_dMag[sInds]
1✔
1060
        WA = TL.int_WA[sInds]
1✔
1061

1062
        # save out file containing photon count info
1063
        if self.record_counts_path is not None and len(self.count_lines) == 0:
1✔
1064
            C_p, C_b, C_sp, C_extra = self.OpticalSystem.Cp_Cb_Csp(
×
1065
                self.TargetList, sInds, fZ, fEZs, dMag, WA, mode, returnExtra=True
1066
            )
1067
            import csv
×
1068

1069
            count_fpath = os.path.join(self.record_counts_path, "counts")
×
1070

1071
            if not os.path.exists(count_fpath):
×
1072
                os.mkdir(count_fpath)
×
1073

1074
            outfile = os.path.join(count_fpath, str(self.seed) + ".csv")
×
1075
            self.count_lines.append(
×
1076
                [
1077
                    "sInd",
1078
                    "HIPs",
1079
                    "C_F0",
1080
                    "C_p0",
1081
                    "C_sr",
1082
                    "C_z",
1083
                    "C_ez",
1084
                    "C_dc",
1085
                    "C_cc",
1086
                    "C_rn",
1087
                    "C_p",
1088
                    "C_b",
1089
                    "C_sp",
1090
                ]
1091
            )
1092

1093
            for i, sInd in enumerate(sInds):
×
1094
                self.count_lines.append(
×
1095
                    [
1096
                        sInd,
1097
                        self.TargetList.Name[sInd],
1098
                        C_extra["C_F0"][0].value,
1099
                        C_extra["C_sr"][i].value,
1100
                        C_extra["C_z"][i].value,
1101
                        C_extra["C_ez"][i].value,
1102
                        C_extra["C_dc"][i].value,
1103
                        C_extra["C_cc"][i].value,
1104
                        C_extra["C_rn"][i].value,
1105
                        C_p[i].value,
1106
                        C_b[i].value,
1107
                        C_sp[i].value,
1108
                    ]
1109
                )
1110

1111
            with open(outfile, "w") as csvfile:
×
1112
                c = csv.writer(csvfile)
×
1113
                c.writerows(self.count_lines)
×
1114

1115
        intTimes = self.OpticalSystem.calc_intTime(
1✔
1116
            self.TargetList, sInds, fZ, fEZs, dMag, WA, mode
1117
        )
1118
        intTimes[~np.isfinite(intTimes)] = 0 * u.d
1✔
1119

1120
        return intTimes
1✔
1121

1122
    def choose_next_target(self, old_sInd, sInds, slewTimes, intTimes):
1✔
1123
        """Helper method for method next_target to simplify alternative implementations.
1124

1125
        Given a subset of targets (pre-filtered by method next_target or some
1126
        other means), select the best next one. The prototype uses completeness
1127
        as the sole heuristic.
1128

1129
        Args:
1130
            old_sInd (int):
1131
                Index of the previous target star
1132
            sInds (~numpy.ndarray(int)):
1133
                Indices of available targets
1134
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1135
                slew times to all stars (must be indexed by sInds)
1136
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1137
                Integration times for detection in units of day
1138

1139
        Returns:
1140
            tuple:
1141
                sInd (int):
1142
                    Index of next target star
1143
                waitTime (:py:class:`~astropy.units.Quantity`):
1144
                    Some strategic amount of time to wait in case an occulter slew is
1145
                    desired (default is None)
1146

1147
        """
1148

1149
        Comp = self.Completeness
1✔
1150
        TL = self.TargetList
1✔
1151
        TK = self.TimeKeeping
1✔
1152
        OS = self.OpticalSystem
1✔
1153
        Obs = self.Observatory
1✔
1154
        allModes = OS.observingModes
1✔
1155

1156
        # cast sInds to array
1157
        sInds = np.array(sInds, ndmin=1, copy=copy_if_needed)
1✔
1158
        # calculate dt since previous observation
1159
        dt = TK.currentTimeNorm.copy() + slewTimes[sInds] - self.lastObsTimes[sInds]
1✔
1160
        # get dynamic completeness values
1161
        comps = Comp.completeness_update(TL, sInds, self.starVisits[sInds], dt)
1✔
1162
        # choose target with maximum completeness
1163
        sInd = np.random.choice(sInds[comps == max(comps)])
1✔
1164

1165
        # Check if exoplanetObsTime would be exceeded
1166
        mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
1167
        (
1✔
1168
            maxIntTimeOBendTime,
1169
            maxIntTimeExoplanetObsTime,
1170
            maxIntTimeMissionLife,
1171
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
1172
        maxIntTime = min(
1✔
1173
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1174
        )  # Maximum intTime allowed
1175
        intTimes2 = self.calc_targ_intTime(
1✔
1176
            np.array([sInd]), TK.currentTimeAbs.copy(), mode
1177
        )
1178
        if (
1✔
1179
            intTimes2 > maxIntTime
1180
        ):  # check if max allowed integration time would be exceeded
1181
            self.vprint("max allowed integration time would be exceeded")
×
1182
            sInd = None
×
1183
            waitTime = 1.0 * u.d
×
1184
            return sInd, waitTime
×
1185

1186
        return (
1✔
1187
            sInd,
1188
            slewTimes[sInd],
1189
        )  # if coronagraph or first sInd, waitTime will be 0 days
1190

1191
    def refineOcculterSlews(self, old_sInd, sInds, slewTimes, obsTimes, sd, mode):
1✔
1192
        """Refines/filters/chooses occulter slews based on time constraints
1193

1194
        Refines the selection of occulter slew times by filtering based on mission time
1195
        constraints and selecting the best slew time for each star. This method calls on
1196
        other occulter methods within SurveySimulation depending on how slew times were
1197
        calculated prior to calling this function (i.e. depending on which
1198
        implementation of the Observatory module is used).
1199

1200
        Args:
1201
            old_sInd (int):
1202
                Index of the previous target star
1203
            sInds (~numpy.ndarray(int)):
1204
                Indices of available targets
1205
            slewTimes (astropy quantity array):
1206
                slew times to all stars (must be indexed by sInds)
1207
            obsTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1208
                A binary array with TargetList.nStars rows and
1209
                (missionFinishAbs-missionStart)/dt columns
1210
                where dt is 1 day by default. A value of 1 indicates the star is in
1211
                keepout (and therefore cannot be observed). A value of 0 indicates the
1212
                star is not in keepout and may be observed.
1213
            sd (~astropy.units.Quantity(~numpy.ndarray(float))):
1214
                Angular separation between stars in rad
1215
            mode (dict):
1216
                Selected observing mode for detection
1217

1218
        Returns:
1219
            tuple:
1220
                sInds (int):
1221
                    Indeces of next target star
1222
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1223
                    slew times to all stars (must be indexed by sInds)
1224
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1225
                    Integration times for detection in units of day
1226
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1227
                    Delta-V used to transfer to new star line of sight in unis of m/s
1228
        """
1229

1230
        Obs = self.Observatory
×
1231
        TL = self.TargetList
×
1232

1233
        # initializing arrays
1234
        obsTimeArray = np.zeros([TL.nStars, 50]) * u.d
×
1235
        intTimeArray = np.zeros([TL.nStars, 2]) * u.d
×
1236

1237
        for n in sInds:
×
1238
            obsTimeArray[n, :] = (
×
1239
                np.linspace(obsTimes[0, n].value, obsTimes[1, n].value, 50) * u.d
1240
            )
1241
        intTimeArray[sInds, 0] = self.calc_targ_intTime(
×
1242
            sInds, Time(obsTimeArray[sInds, 0], format="mjd", scale="tai"), mode
1243
        )
1244
        intTimeArray[sInds, 1] = self.calc_targ_intTime(
×
1245
            sInds, Time(obsTimeArray[sInds, -1], format="mjd", scale="tai"), mode
1246
        )
1247

1248
        # determining which scheme to use to filter slews
1249
        obsModName = Obs.__class__.__name__
×
1250

1251
        # slew times have not been calculated/decided yet (SotoStarshade)
1252
        if obsModName == "SotoStarshade":
×
1253
            sInds, intTimes, slewTimes, dV = self.findAllowableOcculterSlews(
×
1254
                sInds,
1255
                old_sInd,
1256
                sd[sInds],
1257
                slewTimes[sInds],
1258
                obsTimeArray[sInds, :],
1259
                intTimeArray[sInds, :],
1260
                mode,
1261
            )
1262

1263
        # slew times were calculated/decided beforehand (Observatory Prototype)
1264
        else:
1265
            sInds, intTimes, slewTimes = self.filterOcculterSlews(
×
1266
                sInds,
1267
                slewTimes[sInds],
1268
                obsTimeArray[sInds, :],
1269
                intTimeArray[sInds, :],
1270
                mode,
1271
            )
1272
            dV = np.zeros(len(sInds)) * u.m / u.s
×
1273

1274
        return sInds, slewTimes, intTimes, dV
×
1275

1276
    def filterOcculterSlews(self, sInds, slewTimes, obsTimeArray, intTimeArray, mode):
1✔
1277
        """Filters occulter slews that have already been calculated/selected.
1278

1279
        Used by the refineOcculterSlews method when slew times have been selected
1280
        a priori. This method filters out slews that are not within desired observing
1281
        blocks, the maximum allowed integration time, and are outside of future
1282
        keepouts.
1283

1284
        Args:
1285
            sInds (~numpy.ndarray(int)):
1286
                Indices of available targets
1287
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1288
                slew times to all stars (must be indexed by sInds)
1289
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1290
                Array of times during which a star is out of keepout, has shape
1291
                nx50 where n is the number of stars in sInds. Unit of days
1292
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1293
                Array of integration times for each time in obsTimeArray, has shape
1294
                nx2 where n is the number of stars in sInds. Unit of days
1295
            mode (dict):
1296
                Selected observing mode for detection
1297

1298
        Returns:
1299
            tuple:
1300
                sInds (int):
1301
                    Indeces of next target star
1302
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1303
                    Integration times for detection in units of day
1304
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1305
                    slew times to all stars (must be indexed by sInds)
1306
        """
1307

1308
        TK = self.TimeKeeping
×
1309
        Obs = self.Observatory
×
1310

1311
        # allocate settling time + overhead time
1312
        tmpCurrentTimeAbs = (
×
1313
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1314
        )
1315
        tmpCurrentTimeNorm = (
×
1316
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1317
        )
1318

1319
        # 0. lambda function that linearly interpolates Integration Time
1320
        # between obsTimes
1321
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1322
            t - np.array(x[:, 0]).reshape(len(t), 1)
1323
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1324

1325
        # 1. initializing arrays
1326
        obsTimesRange = np.array(
×
1327
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1328
        )  # nx2 array with start and end times of obsTimes for each star
1329
        intTimesRange = np.array([intTimeArray[:, 0], intTimeArray[:, -1]])
×
1330

1331
        OBnumbers = np.zeros(
×
1332
            [len(sInds), 1]
1333
        )  # for each sInd, will show during which OB observations will take place
1334
        maxIntTimes = np.zeros([len(sInds), 1]) * u.d
×
1335

1336
        intTimes = (
×
1337
            linearInterp(
1338
                intTimesRange.T,
1339
                obsTimesRange.T,
1340
                (tmpCurrentTimeAbs + slewTimes).reshape(len(sInds), 1).value,
1341
            )
1342
            * u.d
1343
        )  # calculate intTimes for each slew time
1344

1345
        minObsTimeNorm = (obsTimesRange[0, :] - tmpCurrentTimeAbs.value).reshape(
×
1346
            [len(sInds), 1]
1347
        )
1348
        maxObsTimeNorm = (obsTimesRange[1, :] - tmpCurrentTimeAbs.value).reshape(
×
1349
            [len(sInds), 1]
1350
        )
1351
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1352

1353
        # 2. find OBnumber for each sInd's slew time
1354
        if len(TK.OBendTimes) > 1:
×
1355
            for i in range(len(sInds)):
×
1356
                S = np.where(
×
1357
                    TK.OBstartTimes.value - tmpCurrentTimeNorm.value
1358
                    < slewTimes[i].value
1359
                )[0][-1]
1360
                F = np.where(
×
1361
                    TK.OBendTimes.value - tmpCurrentTimeNorm.value < slewTimes[i].value
1362
                )[0]
1363

1364
                # case when slews are in the first OB
1365
                if F.shape[0] == 0:
×
1366
                    F = -1
×
1367
                else:
1368
                    F = F[-1]
×
1369

1370
                # slew occurs within an OB (nth OB has started but hasn't ended)
1371
                if S != F:
×
1372
                    OBnumbers[i] = S
×
1373
                    (
×
1374
                        maxIntTimeOBendTime,
1375
                        maxIntTimeExoplanetObsTime,
1376
                        maxIntTimeMissionLife,
1377
                    ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, TK.OBstartTimes[S], S)
1378
                    maxIntTimes[i] = min(
×
1379
                        maxIntTimeOBendTime,
1380
                        maxIntTimeExoplanetObsTime,
1381
                        maxIntTimeMissionLife,
1382
                    )  # Maximum intTime allowed
1383

1384
                # slew occurs between OBs, badbadnotgood
1385
                else:
1386
                    OBnumbers[i] = -1
×
1387
                    maxIntTimes[i] = 0 * u.d
×
1388
            OBstartTimeNorm = (
×
1389
                TK.OBstartTimes[np.array(OBnumbers, dtype=int)].value
1390
                - tmpCurrentTimeNorm.value
1391
            )
1392
        else:
1393
            (
×
1394
                maxIntTimeOBendTime,
1395
                maxIntTimeExoplanetObsTime,
1396
                maxIntTimeMissionLife,
1397
            ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm)
1398
            maxIntTimes[:] = min(
×
1399
                maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1400
            )  # Maximum intTime allowed
1401
            OBstartTimeNorm = np.zeros(OBnumbers.shape)
×
1402

1403
        # finding the minimum possible slew time
1404
        # (either OBstart or when star JUST leaves keepout)
1405
        minAllowedSlewTime = np.max([OBstartTimeNorm, minObsTimeNorm], axis=0)
×
1406

1407
        # 3. start filtering process
1408
        good_inds = np.where((OBnumbers >= 0) & (ObsTimeRange > intTimes.value))[0]
×
1409
        # star slew within OB -AND- can finish observing
1410
        # before it goes back into keepout
1411
        if good_inds.shape[0] > 0:
×
1412
            # the good ones
1413
            sInds = sInds[good_inds]
×
1414
            slewTimes = slewTimes[good_inds]
×
1415
            intTimes = intTimes[good_inds]
×
1416
            OBstartTimeNorm = OBstartTimeNorm[good_inds]
×
1417
            minAllowedSlewTime = minAllowedSlewTime[good_inds]
×
1418

1419
            # maximum allowed slew time based on integration times
1420
            maxAllowedSlewTime = maxIntTimes[good_inds].value - intTimes.value
×
1421
            maxAllowedSlewTime[maxAllowedSlewTime < 0] = -np.inf
×
1422
            maxAllowedSlewTime += OBstartTimeNorm  # calculated rel to currentTime norm
×
1423

1424
            # checking to see if slewTimes are allowed
1425
            good_inds = np.where(
×
1426
                (slewTimes.reshape([len(sInds), 1]).value > minAllowedSlewTime)
1427
                & (slewTimes.reshape([len(sInds), 1]).value < maxAllowedSlewTime)
1428
            )[0]
1429

1430
            slewTimes = slewTimes[good_inds]
×
1431
        else:
1432
            slewTimes = slewTimes[good_inds]
×
1433

1434
        return sInds[good_inds], intTimes[good_inds].flatten(), slewTimes
×
1435

1436
    def findAllowableOcculterSlews(
1✔
1437
        self, sInds, old_sInd, sd, slewTimes, obsTimeArray, intTimeArray, mode
1438
    ):
1439
        """Finds an array of allowable slew times for each star
1440

1441
        Used by the refineOcculterSlews method when slew times have NOT been selected
1442
        a priori. This method creates nx50 arrays (where the row corresponds to a
1443
        specific star and the column corresponds to a future point in time relative to
1444
        currentTime).
1445

1446
        These arrays are initially zero but are populated with the corresponding values
1447
        (slews, intTimes, etc) if slewing to that time point (i.e. beginning an
1448
        observation) would lead to a successful observation. A "successful observation"
1449
        is defined by certain conditions relating to keepout and the komap, observing
1450
        blocks, mission lifetime, and some constraints on the dVmap calculation in
1451
        SotoStarshade. Each star will likely have a range of slewTimes that would lead
1452
        to a successful observation -- another method is then called to select the best
1453
        of these slewTimes.
1454

1455
        Args:
1456
            sInds (~numpy.ndarray(int)):
1457
                Indices of available targets
1458
            old_sInd (int):
1459
                Index of the previous target star
1460
            sd (~astropy.units.Quantity(~numpy.ndarray(float))):
1461
                Angular separation between stars in rad
1462
            slewTimes (astropy quantity array):
1463
                slew times to all stars (must be indexed by sInds)
1464
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1465
                Array of times during which a star is out of keepout, has shape
1466
                nx50 where n is the number of stars in sInds
1467
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1468
                Array of integration times for each time in obsTimeArray, has shape
1469
                nx50 where n is the number of stars in sInds
1470
            mode (dict):
1471
                Selected observing mode for detection
1472

1473
        Returns:
1474
            tuple:
1475
                sInds (numpy.ndarray(int)):
1476
                    Indices of next target star
1477
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1478
                    slew times to all stars (must be indexed by sInds)
1479
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1480
                    Integration times for detection in units of day
1481
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1482
                    Delta-V used to transfer to new star line of sight in unis of m/s
1483
        """
1484
        TK = self.TimeKeeping
×
1485
        Obs = self.Observatory
×
1486
        TL = self.TargetList
×
1487

1488
        # 0. lambda function that linearly interpolates Integration
1489
        # Time between obsTimes
1490
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1491
            t - np.array(x[:, 0]).reshape(len(t), 1)
1492
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1493

1494
        # allocate settling time + overhead time
1495
        tmpCurrentTimeAbs = (
×
1496
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1497
        )
1498
        tmpCurrentTimeNorm = (
×
1499
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1500
        )
1501

1502
        # 1. initializing arrays
1503
        obsTimes = np.array(
×
1504
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1505
        )  # nx2 array with start and end times of obsTimes for each star
1506
        intTimes_int = (
×
1507
            np.zeros(obsTimeArray.shape) * u.d
1508
        )  # initializing intTimes of shape nx50 then interpolating
1509
        intTimes_int = (
×
1510
            np.hstack(
1511
                [
1512
                    intTimeArray[:, 0].reshape(len(sInds), 1).value,
1513
                    linearInterp(
1514
                        intTimeArray.value, obsTimes.T, obsTimeArray[:, 1:].value
1515
                    ),
1516
                ]
1517
            )
1518
            * u.d
1519
        )
1520
        allowedSlewTimes = np.zeros(obsTimeArray.shape) * u.d
×
1521
        allowedintTimes = np.zeros(obsTimeArray.shape) * u.d
×
1522
        allowedCharTimes = np.zeros(obsTimeArray.shape) * u.d
×
1523
        obsTimeArrayNorm = obsTimeArray.value - tmpCurrentTimeAbs.value
×
1524

1525
        # obsTimes -> relative to current Time
1526
        try:
×
1527
            minObsTimeNorm = np.array([np.min(v[v > 0]) for v in obsTimeArrayNorm])
×
1528
        except:  # noqa: E722
×
1529
            # an error pops up sometimes at the end of the mission, this fixes it
1530
            # TODO: define the error type that occurs
1531
            # rewrite to avoid a try/except if possible
1532
            minObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1533

1534
        maxObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1535
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1536

1537
        # getting max possible intTime
1538
        (
×
1539
            maxIntTimeOBendTime,
1540
            maxIntTimeExoplanetObsTime,
1541
            maxIntTimeMissionLife,
1542
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm, TK.OBnumber)
1543
        maxIntTime = min(
×
1544
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1545
        )  # Maximum intTime allowed
1546

1547
        # 2. giant array of min and max slew times, starts at current time, ends when
1548
        # stars enter keepout (all same size). Each entry either has a slew time value
1549
        # if a slew is allowed at that date or 0 if slewing is not allowed
1550

1551
        # first filled in for the current OB
1552
        minAllowedSlewTimes = np.array(
×
1553
            [minObsTimeNorm.T] * len(intTimes_int.T)
1554
        ).T  # just to make it nx50 so it plays nice with the other arrays
1555
        maxAllowedSlewTimes = maxIntTime.value - intTimes_int.value
×
1556
        maxAllowedSlewTimes[maxAllowedSlewTimes > Obs.occ_dtmax.value] = (
×
1557
            Obs.occ_dtmax.value
1558
        )
1559

1560
        # conditions that must be met to define an allowable slew time
1561
        cond1 = (
×
1562
            minAllowedSlewTimes >= Obs.occ_dtmin.value
1563
        )  # minimum dt time in dV map interpolant
1564
        cond2 = (
×
1565
            maxAllowedSlewTimes <= Obs.occ_dtmax.value
1566
        )  # maximum dt time in dV map interpolant
1567
        cond3 = maxAllowedSlewTimes > minAllowedSlewTimes
×
1568
        cond4 = intTimes_int.value < ObsTimeRange.reshape(len(sInds), 1)
×
1569

1570
        conds = cond1 & cond2 & cond3 & cond4
×
1571
        minAllowedSlewTimes[np.invert(conds)] = (
×
1572
            np.inf
1573
        )  # these are filtered during the next filter
1574
        maxAllowedSlewTimes[np.invert(conds)] = -np.inf
×
1575

1576
        # one last condition to meet
1577
        map_i, map_j = np.where(
×
1578
            (obsTimeArrayNorm > minAllowedSlewTimes)
1579
            & (obsTimeArrayNorm < maxAllowedSlewTimes)
1580
        )
1581

1582
        # 2.5 if any stars are slew-able to within this OB block, populate
1583
        # "allowedSlewTimes", a running tally of possible slews
1584
        # within the time range a star is observable (out of keepout)
1585
        if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1586
            allowedSlewTimes[map_i, map_j] = obsTimeArrayNorm[map_i, map_j] * u.d
×
1587
            allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1588
            allowedCharTimes[map_i, map_j] = maxIntTime - intTimes_int[map_i, map_j]
×
1589

1590
        # 3. search future OBs
1591
        OB_withObsStars = (
×
1592
            TK.OBstartTimes.value - np.min(obsTimeArrayNorm) - tmpCurrentTimeNorm.value
1593
        )  # OBs within which any star is observable
1594

1595
        if any(OB_withObsStars > 0):
×
1596
            nOBstart = np.argmin(np.abs(OB_withObsStars))
×
1597
            nOBend = np.argmax(OB_withObsStars)
×
1598

1599
            # loop through the next 5 OBs (or until mission is over if there are less
1600
            # than 5 OBs in the future)
1601
            for i in np.arange(nOBstart, np.min([nOBend, nOBstart + 5])):
×
1602

1603
                # max int Times for the next OB
1604
                (
×
1605
                    maxIntTimeOBendTime,
1606
                    maxIntTimeExoplanetObsTime,
1607
                    maxIntTimeMissionLife,
1608
                ) = TK.get_ObsDetectionMaxIntTime(
1609
                    Obs, mode, TK.OBstartTimes[i + 1], i + 1
1610
                )
1611
                maxIntTime_nOB = min(
×
1612
                    maxIntTimeOBendTime,
1613
                    maxIntTimeExoplanetObsTime,
1614
                    maxIntTimeMissionLife,
1615
                )  # Maximum intTime allowed
1616

1617
                # min and max slew times rel to current Time (norm)
1618
                nOBstartTimeNorm = np.array(
×
1619
                    [TK.OBstartTimes[i + 1].value - tmpCurrentTimeNorm.value]
1620
                    * len(sInds)
1621
                )
1622

1623
                # min slew times for stars start either whenever the star first leaves
1624
                # keepout or when next OB stars, whichever comes last
1625
                minAllowedSlewTimes_nOB = np.array(
×
1626
                    [np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T]
1627
                    * len(maxAllowedSlewTimes.T)
1628
                ).T
1629
                maxAllowedSlewTimes_nOB = (
×
1630
                    nOBstartTimeNorm.reshape(len(sInds), 1)
1631
                    + maxIntTime_nOB.value
1632
                    - intTimes_int.value
1633
                )
1634
                maxAllowedSlewTimes_nOB[
×
1635
                    maxAllowedSlewTimes_nOB > Obs.occ_dtmax.value
1636
                ] = Obs.occ_dtmax.value
1637

1638
                # amount of time left when the stars are still out of keepout
1639
                ObsTimeRange_nOB = (
×
1640
                    maxObsTimeNorm
1641
                    - np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T
1642
                )
1643

1644
                # condition to be met for an allowable slew time
1645
                cond1 = minAllowedSlewTimes_nOB >= Obs.occ_dtmin.value
×
1646
                cond2 = maxAllowedSlewTimes_nOB <= Obs.occ_dtmax.value
×
1647
                cond3 = maxAllowedSlewTimes_nOB > minAllowedSlewTimes_nOB
×
1648
                cond4 = intTimes_int.value < ObsTimeRange_nOB.reshape(len(sInds), 1)
×
1649
                cond5 = intTimes_int.value < maxIntTime_nOB.value
×
1650
                conds = cond1 & cond2 & cond3 & cond4 & cond5
×
1651

1652
                minAllowedSlewTimes_nOB[np.invert(conds)] = np.inf
×
1653
                maxAllowedSlewTimes_nOB[np.invert(conds)] = -np.inf
×
1654

1655
                # one last condition
1656
                map_i, map_j = np.where(
×
1657
                    (obsTimeArrayNorm > minAllowedSlewTimes_nOB)
1658
                    & (obsTimeArrayNorm < maxAllowedSlewTimes_nOB)
1659
                )
1660

1661
                # 3.33 populate the running tally of allowable slew times if it meets
1662
                # all conditions
1663
                if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1664
                    allowedSlewTimes[map_i, map_j] = (
×
1665
                        obsTimeArrayNorm[map_i, map_j] * u.d
1666
                    )
1667
                    allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1668
                    allowedCharTimes[map_i, map_j] = (
×
1669
                        maxIntTime_nOB - intTimes_int[map_i, map_j]
1670
                    )
1671

1672
        # 3.67 filter out any stars that are not observable at all
1673
        filterDuds = np.sum(allowedSlewTimes, axis=1) > 0.0
×
1674
        sInds = sInds[filterDuds]
×
1675

1676
        # 4. choose a slew time for each available star
1677
        # calculate dVs for each possible slew time for each star
1678
        allowed_dVs = Obs.calculate_dV(
×
1679
            TL,
1680
            old_sInd,
1681
            sInds,
1682
            sd[filterDuds],
1683
            allowedSlewTimes[filterDuds, :],
1684
            tmpCurrentTimeAbs,
1685
        )
1686

1687
        if len(sInds.tolist()) > 0:
×
1688
            # select slew time for each star
1689
            dV_inds = np.arange(0, len(sInds))
×
1690
            sInds, intTime, slewTime, dV = self.chooseOcculterSlewTimes(
×
1691
                sInds,
1692
                allowedSlewTimes[filterDuds, :],
1693
                allowed_dVs[dV_inds, :],
1694
                allowedintTimes[filterDuds, :],
1695
                allowedCharTimes[filterDuds, :],
1696
            )
1697

1698
            return sInds, intTime, slewTime, dV
×
1699

1700
        else:
1701
            empty = np.asarray([], dtype=int)
×
1702
            return empty, empty * u.d, empty * u.d, empty * u.m / u.s
×
1703

1704
    def chooseOcculterSlewTimes(self, sInds, slewTimes, dV, intTimes, charTimes):
1✔
1705
        """Selects the best slew time for each star
1706

1707
        This method searches through an array of permissible slew times for
1708
        each star and chooses the best slew time for the occulter based on
1709
        maximizing possible characterization time for that particular star (as
1710
        a default).
1711

1712
        Args:
1713
            sInds (~numpy.ndarray(int)):
1714
                Indices of available targets
1715
            slewTimes (astropy quantity array):
1716
                slew times to all stars (must be indexed by sInds)
1717
            dV (~astropy.units.Quantity(~numpy.ndarray(float))):
1718
                Delta-V used to transfer to new star line of sight in unis of m/s
1719
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1720
                Integration times for detection in units of day
1721
            charTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1722
                Time left over after integration which could be used for
1723
                characterization in units of day
1724

1725
        Returns:
1726
            tuple:
1727
                sInds (int):
1728
                    Indeces of next target star
1729
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1730
                    slew times to all stars (must be indexed by sInds)
1731
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1732
                    Integration times for detection in units of day
1733
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1734
                    Delta-V used to transfer to new star line of sight in unis of m/s
1735
        """
1736

1737
        # selection criteria for each star slew
1738
        good_j = np.argmax(
×
1739
            charTimes, axis=1
1740
        )  # maximum possible characterization time available
1741
        good_i = np.arange(0, len(sInds))
×
1742

1743
        dV = dV[good_i, good_j]
×
1744
        intTime = intTimes[good_i, good_j]
×
1745
        slewTime = slewTimes[good_i, good_j]
×
1746

1747
        return sInds, intTime, slewTime, dV
×
1748

1749
    def observation_detection(self, sInd, intTime, mode):
1✔
1750
        """Determines SNR and detection status for a given integration time
1751
        for detection. Also updates the lastDetected and starRevisit lists.
1752

1753
        Args:
1754
            sInd (int):
1755
                Integer index of the star of interest
1756
            intTime (~astropy.units.Quantity(~numpy.ndarray(float))):
1757
                Selected star integration time for detection in units of day.
1758
                Defaults to None.
1759
            mode (dict):
1760
                Selected observing mode for detection
1761

1762
        Returns:
1763
            tuple:
1764
                detected (numpy.ndarray(int)):
1765
                    Detection status for each planet orbiting the observed target star:
1766
                    1 is detection, 0 missed detection, -1 below IWA, and -2 beyond OWA
1767
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
1768
                    Surface brightness of local zodiacal light in units of 1/arcsec2
1769
                systemParams (dict):
1770
                    Dictionary of time-dependant planet properties averaged over the
1771
                    duration of the integration
1772
                SNR (numpy.darray(float)):
1773
                    Detection signal-to-noise ratio of the observable planets
1774
                FA (bool):
1775
                    False alarm (false positive) boolean
1776

1777
        """
1778

1779
        PPop = self.PlanetPopulation
1✔
1780
        ZL = self.ZodiacalLight
1✔
1781
        PPro = self.PostProcessing
1✔
1782
        TL = self.TargetList
1✔
1783
        SU = self.SimulatedUniverse
1✔
1784
        Obs = self.Observatory
1✔
1785
        TK = self.TimeKeeping
1✔
1786

1787
        # Save Current Time before attempting time allocation
1788
        currentTimeNorm = TK.currentTimeNorm.copy()
1✔
1789
        currentTimeAbs = TK.currentTimeAbs.copy()
1✔
1790

1791
        # Allocate Time
1792
        extraTime = intTime * (mode["timeMultiplier"] - 1.0)  # calculates extraTime
1✔
1793
        success = TK.allocate_time(
1✔
1794
            intTime + extraTime + Obs.settlingTime + mode["syst"]["ohTime"], True
1795
        )
1796
        assert success, "Could not allocate observation detection time ({}).".format(
1✔
1797
            intTime + extraTime + Obs.settlingTime + mode["syst"]["ohTime"]
1798
        )
1799
        # calculates partial time to be added for every ntFlux
1800
        dt = intTime / float(self.ntFlux)
1✔
1801
        # find indices of planets around the target
1802
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
1803

1804
        # initialize outputs
1805
        detected = np.array([], dtype=int)
1✔
1806
        fZ = 0.0 / u.arcsec**2
1✔
1807
        # write current system params by default
1808
        systemParams = SU.dump_system_params(sInd)
1✔
1809
        SNR = np.zeros(len(pInds))
1✔
1810

1811
        # if any planet, calculate SNR
1812
        if len(pInds) > 0:
1✔
1813
            # initialize arrays for SNR integration
1814
            fZs = np.zeros(self.ntFlux) / u.arcsec**2
1✔
1815
            systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
1816
            Ss = np.zeros((self.ntFlux, len(pInds)))
1✔
1817
            Ns = np.zeros((self.ntFlux, len(pInds)))
1✔
1818
            # accounts for the time since the current time
1819
            timePlus = Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
1✔
1820
            # integrate the signal (planet flux) and noise
1821
            for i in range(self.ntFlux):
1✔
1822
                # allocate first half of dt
1823
                timePlus += dt / 2.0
1✔
1824
                # calculate current zodiacal light brightness
1825
                fZs[i] = ZL.fZ(Obs, TL, sInd, currentTimeAbs + timePlus, mode)[0]
1✔
1826
                # propagate the system to match up with current time
1827
                SU.propag_system(
1✔
1828
                    sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
1829
                )
1830
                self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
1831
                # save planet parameters
1832
                systemParamss[i] = SU.dump_system_params(sInd)
1✔
1833
                # calculate signal and noise (electron count rates)
1834
                Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
1835
                    sInd, pInds, dt, mode, fZ=fZs[i]
1836
                )
1837
                # allocate second half of dt
1838
                timePlus += dt / 2.0
1✔
1839

1840
            # average output parameters
1841
            fZ = np.mean(fZs)
1✔
1842
            systemParams = {
1✔
1843
                key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
1844
                / float(self.ntFlux)
1845
                for key in sorted(systemParamss[0])
1846
            }
1847
            # calculate SNR
1848
            S = Ss.sum(0)
1✔
1849
            N = Ns.sum(0)
1✔
1850
            SNR[N > 0] = S[N > 0] / N[N > 0]
1✔
1851

1852
        # if no planet, just save zodiacal brightness in the middle of the integration
1853
        else:
1854
            totTime = intTime * (mode["timeMultiplier"])
1✔
1855
            fZ = ZL.fZ(Obs, TL, sInd, currentTimeAbs + totTime / 2.0, mode)[0]
1✔
1856

1857
        # find out if a false positive (false alarm) or any false negative
1858
        # (missed detections) have occurred
1859
        FA, MD = PPro.det_occur(SNR, mode, TL, sInd, intTime)
1✔
1860

1861
        # populate detection status array
1862
        # 1:detected, 0:missed, -1:below IWA, -2:beyond OWA
1863
        if len(pInds) > 0:
1✔
1864
            detected = (~MD).astype(int)
1✔
1865
            WA = (
1✔
1866
                np.array(
1867
                    [
1868
                        systemParamss[x]["WA"].to("arcsec").value
1869
                        for x in range(len(systemParamss))
1870
                    ]
1871
                )
1872
                * u.arcsec
1873
            )
1874
            detected[np.all(WA < mode["IWA"], 0)] = -1
1✔
1875
            detected[np.all(WA > mode["OWA"], 0)] = -2
1✔
1876

1877
        # if planets are detected, calculate the minimum apparent separation
1878
        smin = None
1✔
1879
        det = detected == 1  # If any of the planets around the star have been detected
1✔
1880
        if np.any(det):
1✔
1881
            smin = np.min(SU.s[pInds[det]])
1✔
1882
            log_det = "   - Detected planet inds %s (%s/%s)" % (
1✔
1883
                pInds[det],
1884
                len(pInds[det]),
1885
                len(pInds),
1886
            )
1887
            self.logger.info(log_det)
1✔
1888
            self.vprint(log_det)
1✔
1889

1890
        # populate the lastDetected array by storing det, fEZ, dMag, and WA
1891
        self.lastDetected[sInd, :] = [
1✔
1892
            det,
1893
            systemParams["fEZ"].to("1/arcsec2").value,
1894
            systemParams["dMag"],
1895
            systemParams["WA"].to("arcsec").value,
1896
        ]
1897

1898
        # in case of a FA, generate a random delta mag (between PPro.FAdMag0 and
1899
        # TL.intCutoff_dMag) and working angle (between IWA and min(OWA, a_max))
1900
        if FA:
1✔
1901
            WA = (
×
1902
                np.random.uniform(
1903
                    mode["IWA"].to("arcsec").value,
1904
                    np.minimum(mode["OWA"], np.arctan(max(PPop.arange) / TL.dist[sInd]))
1905
                    .to("arcsec")
1906
                    .value,
1907
                )
1908
                * u.arcsec
1909
            )
1910
            dMag = np.random.uniform(PPro.FAdMag0(WA), TL.intCutoff_dMag)
×
1911
            self.lastDetected[sInd, 0] = np.append(self.lastDetected[sInd, 0], True)
×
1912
            self.lastDetected[sInd, 1] = np.append(
×
1913
                self.lastDetected[sInd, 1], ZL.fEZ0.to("1/arcsec2").value
1914
            )
1915
            self.lastDetected[sInd, 2] = np.append(self.lastDetected[sInd, 2], dMag)
×
1916
            self.lastDetected[sInd, 3] = np.append(
×
1917
                self.lastDetected[sInd, 3], WA.to("arcsec").value
1918
            )
1919
            sminFA = np.tan(WA) * TL.dist[sInd].to("AU")
×
1920
            smin = np.minimum(smin, sminFA) if smin is not None else sminFA
×
1921
            log_FA = "   - False Alarm (WA=%s, dMag=%s)" % (
×
1922
                np.round(WA, 3),
1923
                np.round(dMag, 1),
1924
            )
1925
            self.logger.info(log_FA)
×
1926
            self.vprint(log_FA)
×
1927

1928
        # Schedule Target Revisit
1929
        self.scheduleRevisit(sInd, smin, det, pInds)
1✔
1930

1931
        if self.make_debug_bird_plots:
1✔
1932
            from tools.obs_plot import obs_plot
×
1933

1934
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, detected)
×
1935

1936
        return detected.astype(int), fZ, systemParams, SNR, FA
1✔
1937

1938
    def scheduleRevisit(self, sInd, smin, det, pInds):
1✔
1939
        """A Helper Method for scheduling revisits after observation detection
1940

1941
        Updates self.starRevisit attribute only
1942

1943
        Args:
1944
            sInd (int):
1945
                sInd of the star just detected
1946
            smin (~astropy.units.Quantity):
1947
                minimum separation of the planet to star of planet just detected
1948
            det (~np.ndarray(bool)):
1949
                Detection status of all planets in target system
1950
            pInds (~np.ndarray(int)):
1951
                Indices of planets in the target system
1952

1953
        Returns:
1954
            None
1955

1956
        """
1957
        TK = self.TimeKeeping
1✔
1958
        TL = self.TargetList
1✔
1959
        SU = self.SimulatedUniverse
1✔
1960

1961
        # in both cases (detection or false alarm), schedule a revisit
1962
        # based on minimum separation
1963
        Ms = TL.MsTrue[sInd]
1✔
1964
        if smin is not None:  # smin is None if no planet was detected
1✔
1965
            sp = smin
1✔
1966
            if np.any(det):
1✔
1967
                pInd_smin = pInds[det][np.argmin(SU.s[pInds[det]])]
1✔
1968
                Mp = SU.Mp[pInd_smin]
1✔
1969
            else:
1970
                Mp = SU.Mp.mean()
×
1971
            mu = const.G * (Mp + Ms)
1✔
1972
            T = 2.0 * np.pi * np.sqrt(sp**3.0 / mu)
1✔
1973
            t_rev = TK.currentTimeNorm.copy() + T / 2.0
1✔
1974
        # otherwise, revisit based on average of population semi-major axis and mass
1975
        else:
1976
            sp = SU.s.mean()
1✔
1977
            Mp = SU.Mp.mean()
1✔
1978
            mu = const.G * (Mp + Ms)
1✔
1979
            T = 2.0 * np.pi * np.sqrt(sp**3.0 / mu)
1✔
1980
            t_rev = TK.currentTimeNorm.copy() + 0.75 * T
1✔
1981

1982
        if self.revisit_wait is not None:
1✔
1983
            t_rev = TK.currentTimeNorm.copy() + self.revisit_wait
×
1984
        # finally, populate the revisit list (NOTE: sInd becomes a float)
1985
        revisit = np.array([sInd, t_rev.to("day").value])
1✔
1986
        if self.starRevisit.size == 0:  # If starRevisit has nothing in it
1✔
1987
            self.starRevisit = np.array([revisit])  # initialize sterRevisit
1✔
1988
        else:
1989
            revInd = np.where(self.starRevisit[:, 0] == sInd)[
1✔
1990
                0
1991
            ]  # indices of the first column of the starRevisit list containing sInd
1992
            if revInd.size == 0:
1✔
1993
                self.starRevisit = np.vstack((self.starRevisit, revisit))
1✔
1994
            else:
1995
                self.starRevisit[revInd, 1] = revisit[1]
×
1996

1997
    def observation_characterization(self, sInd, mode):
1✔
1998
        """Finds if characterizations are possible and relevant information
1999

2000
        Args:
2001
            sInd (int):
2002
                Integer index of the star of interest
2003
            mode (dict):
2004
                Selected observing mode for characterization
2005

2006
        Returns:
2007
            tuple:
2008
                characterized (list(int)):
2009
                    Characterization status for each planet orbiting the observed
2010
                    target star including False Alarm if any, where 1 is full spectrum,
2011
                    -1 partial spectrum, and 0 not characterized
2012
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
2013
                    Surface brightness of local zodiacal light in units of 1/arcsec2
2014
                systemParams (dict):
2015
                    Dictionary of time-dependant planet properties averaged over the
2016
                    duration of the integration
2017
                SNR (numpy.ndarray(float)):
2018
                    Characterization signal-to-noise ratio of the observable planets.
2019
                    Defaults to None.
2020
                intTime (astropy.units.Quantity(numpy.ndarray(float))):
2021
                    Selected star characterization time in units of day.
2022
                    Defaults to None.
2023

2024
        """
2025

2026
        OS = self.OpticalSystem
1✔
2027
        ZL = self.ZodiacalLight
1✔
2028
        TL = self.TargetList
1✔
2029
        SU = self.SimulatedUniverse
1✔
2030
        Obs = self.Observatory
1✔
2031
        TK = self.TimeKeeping
1✔
2032

2033
        # selecting appropriate koMap
2034
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
2035

2036
        # find indices of planets around the target
2037
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2038

2039
        # get the detected status, and check if there was a FA
2040
        det = self.lastDetected[sInd, 0]
1✔
2041
        FA = len(det) == len(pInds) + 1
1✔
2042
        if FA:
1✔
2043
            pIndsDet = np.append(pInds, -1)[det]
×
2044
        else:
2045
            pIndsDet = pInds[det]
1✔
2046

2047
        # initialize outputs, and check if there's anything (planet or FA)
2048
        # to characterize
2049
        characterized = np.zeros(len(det), dtype=int)
1✔
2050
        fZ = 0.0 / u.arcsec**2.0
1✔
2051
        # write current system params by default
2052
        systemParams = SU.dump_system_params(sInd)
1✔
2053
        SNR = np.zeros(len(det))
1✔
2054
        intTime = None
1✔
2055
        if len(det) == 0:  # nothing to characterize
1✔
2056
            return characterized, fZ, systemParams, SNR, intTime
1✔
2057

2058
        # look for last detected planets that have not been fully characterized
2059
        if not (FA):  # only true planets, no FA
1✔
2060
            tochar = self.fullSpectra[pIndsDet] == 0
1✔
2061
        else:  # mix of planets and a FA
2062
            truePlans = pIndsDet[:-1]
×
2063
            tochar = np.append((self.fullSpectra[truePlans] == 0), True)
×
2064

2065
        # 1/ find spacecraft orbital START position including overhead time,
2066
        # and check keepout angle
2067
        if np.any(tochar):
1✔
2068
            # start times
2069
            startTime = (
1✔
2070
                TK.currentTimeAbs.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2071
            )
2072

2073
            # if we're beyond mission end, bail out
2074
            if startTime >= TK.missionFinishAbs:
1✔
2075
                return characterized, fZ, systemParams, SNR, intTime
×
2076

2077
            startTimeNorm = (
1✔
2078
                TK.currentTimeNorm.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2079
            )
2080
            # planets to characterize
2081
            koTimeInd = np.where(np.round(startTime.value) - self.koTimes.value == 0)[
1✔
2082
                0
2083
            ][
2084
                0
2085
            ]  # find indice where koTime is startTime[0]
2086
            # wherever koMap is 1, the target is observable
2087
            tochar[tochar] = koMap[sInd][koTimeInd]
1✔
2088

2089
        # 2/ if any planet to characterize, find the characterization times
2090
        # at the detected fEZ, dMag, and WA
2091
        if np.any(tochar):
1✔
2092
            fZ = ZL.fZ(Obs, TL, [sInd], startTime, mode)
1✔
2093
            fEZ = self.lastDetected[sInd, 1][det][tochar] / u.arcsec**2
1✔
2094
            dMag = self.lastDetected[sInd, 2][det][tochar]
1✔
2095
            WA = self.lastDetected[sInd, 3][det][tochar] * u.arcsec
1✔
2096
            intTimes = np.zeros(len(tochar)) * u.day
1✔
2097
            intTimes[tochar] = OS.calc_intTime(TL, sInd, fZ, fEZ, dMag, WA, mode, TK=TK)
1✔
2098
            intTimes[~np.isfinite(intTimes)] = 0 * u.d
1✔
2099
            # add a predetermined margin to the integration times
2100
            intTimes = intTimes * (1.0 + self.charMargin)
1✔
2101
            # apply time multiplier
2102
            totTimes = intTimes * (mode["timeMultiplier"])
1✔
2103

2104
            # Filter totTimes to make nan integration times correspond to the
2105
            # maximum float value because Time cannot handle nan values
2106
            totTimes[np.where(np.isnan(totTimes))[0]] = np.finfo(np.float64).max * u.d
1✔
2107

2108
            # end times
2109
            endTimes = startTime + totTimes
1✔
2110
            endTimesNorm = startTimeNorm + totTimes
1✔
2111
            # planets to characterize
2112
            tochar = (
1✔
2113
                (totTimes > 0)
2114
                & (totTimes <= OS.intCutoff)
2115
                & (endTimesNorm <= TK.OBendTimes[TK.OBnumber])
2116
            )
2117

2118
        # 3/ is target still observable at the end of any char time?
2119
        if np.any(tochar) and Obs.checkKeepoutEnd:
1✔
2120
            koTimeInds = np.zeros(len(endTimes.value[tochar]), dtype=int)
1✔
2121
            # find index in koMap where each endTime is closest to koTimes
2122
            for t, endTime in enumerate(endTimes.value[tochar]):
1✔
2123
                if endTime > self.koTimes.value[-1]:
1✔
2124
                    # case where endTime exceeds largest koTimes element
2125
                    endTimeInBounds = np.where(
×
2126
                        np.floor(endTime) - self.koTimes.value == 0
2127
                    )[0]
2128
                    koTimeInds[t] = (
×
2129
                        endTimeInBounds[0] if endTimeInBounds.size != 0 else -1
2130
                    )
2131
                else:
2132
                    koTimeInds[t] = np.where(
1✔
2133
                        np.round(endTime) - self.koTimes.value == 0
2134
                    )[0][
2135
                        0
2136
                    ]  # find indice where koTime is endTimes[0]
2137
            tochar[tochar] = [koMap[sInd][koT] if koT >= 0 else 0 for koT in koTimeInds]
1✔
2138

2139
        # 4/ if yes, allocate the overhead time, and perform the characterization
2140
        # for the maximum char time
2141
        if np.any(tochar):
1✔
2142
            # Save Current Time before attempting time allocation
2143
            currentTimeNorm = TK.currentTimeNorm.copy()
1✔
2144
            currentTimeAbs = TK.currentTimeAbs.copy()
1✔
2145

2146
            # Allocate Time
2147
            intTime = np.max(intTimes[tochar])
1✔
2148
            extraTime = intTime * (mode["timeMultiplier"] - 1.0)
1✔
2149
            success = TK.allocate_time(
1✔
2150
                intTime + extraTime + mode["syst"]["ohTime"] + Obs.settlingTime, True
2151
            )  # allocates time
2152
            if not (success):  # Time was not successfully allocated
1✔
2153
                char_intTime = None
×
2154
                lenChar = len(pInds) + 1 if FA else len(pInds)
×
2155
                characterized = np.zeros(lenChar, dtype=float)
×
2156
                char_SNR = np.zeros(lenChar, dtype=float)
×
2157
                char_fZ = 0.0 / u.arcsec**2
×
2158
                char_systemParams = SU.dump_system_params(sInd)
×
2159
                return characterized, char_fZ, char_systemParams, char_SNR, char_intTime
×
2160

2161
            pIndsChar = pIndsDet[tochar]
1✔
2162
            log_char = "   - Charact. planet inds %s (%s/%s detected)" % (
1✔
2163
                pIndsChar,
2164
                len(pIndsChar),
2165
                len(pIndsDet),
2166
            )
2167
            self.logger.info(log_char)
1✔
2168
            self.vprint(log_char)
1✔
2169

2170
            # SNR CALCULATION:
2171
            # first, calculate SNR for observable planets (without false alarm)
2172
            planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
1✔
2173
            SNRplans = np.zeros(len(planinds))
1✔
2174
            if len(planinds) > 0:
1✔
2175
                # initialize arrays for SNR integration
2176
                fZs = np.zeros(self.ntFlux) / u.arcsec**2.0
1✔
2177
                systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
2178
                Ss = np.zeros((self.ntFlux, len(planinds)))
1✔
2179
                Ns = np.zeros((self.ntFlux, len(planinds)))
1✔
2180
                # integrate the signal (planet flux) and noise
2181
                dt = intTime / float(self.ntFlux)
1✔
2182
                timePlus = (
1✔
2183
                    Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2184
                )  # accounts for the time since the current time
2185
                for i in range(self.ntFlux):
1✔
2186
                    # allocate first half of dt
2187
                    timePlus += dt / 2.0
1✔
2188
                    # calculate current zodiacal light brightness
2189
                    fZs[i] = ZL.fZ(Obs, TL, sInd, currentTimeAbs + timePlus, mode)[0]
1✔
2190
                    # propagate the system to match up with current time
2191
                    SU.propag_system(
1✔
2192
                        sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2193
                    )
2194
                    self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
2195
                    # save planet parameters
2196
                    systemParamss[i] = SU.dump_system_params(sInd)
1✔
2197
                    # calculate signal and noise (electron count rates)
2198
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
2199
                        sInd, planinds, dt, mode, fZ=fZs[i]
2200
                    )
2201
                    # allocate second half of dt
2202
                    timePlus += dt / 2.0
1✔
2203

2204
                # average output parameters
2205
                fZ = np.mean(fZs)
1✔
2206
                systemParams = {
1✔
2207
                    key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2208
                    / float(self.ntFlux)
2209
                    for key in sorted(systemParamss[0])
2210
                }
2211
                # calculate planets SNR
2212
                S = Ss.sum(0)
1✔
2213
                N = Ns.sum(0)
1✔
2214
                SNRplans[N > 0] = S[N > 0] / N[N > 0]
1✔
2215
                # allocate extra time for timeMultiplier
2216

2217
            # if only a FA, just save zodiacal brightness in the middle of the
2218
            # integration
2219
            else:
2220
                totTime = intTime * (mode["timeMultiplier"])
×
2221
                fZ = ZL.fZ(Obs, TL, sInd, currentTimeAbs + totTime / 2.0, mode)[0]
×
2222

2223
            # calculate the false alarm SNR (if any)
2224
            SNRfa = []
1✔
2225
            if pIndsChar[-1] == -1:
1✔
2226
                fEZ = self.lastDetected[sInd, 1][-1] / u.arcsec**2.0
×
2227
                dMag = self.lastDetected[sInd, 2][-1]
×
2228
                WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
2229
                C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, fEZ, dMag, WA, mode, TK=TK)
×
2230
                S = (C_p * intTime).decompose().value
×
2231
                N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2232
                SNRfa = S / N if N > 0.0 else 0.0
×
2233

2234
            # save all SNRs (planets and FA) to one array
2235
            SNRinds = np.where(det)[0][tochar]
1✔
2236
            SNR[SNRinds] = np.append(SNRplans, SNRfa)
1✔
2237

2238
            # now, store characterization status: 1 for full spectrum,
2239
            # -1 for partial spectrum, 0 for not characterized
2240
            char = SNR >= mode["SNR"]
1✔
2241
            # initialize with full spectra
2242
            characterized = char.astype(int)
1✔
2243
            WAchar = self.lastDetected[sInd, 3][char] * u.arcsec
1✔
2244
            # find the current WAs of characterized planets
2245
            WAs = systemParams["WA"]
1✔
2246
            if FA:
1✔
2247
                WAs = np.append(WAs, self.lastDetected[sInd, 3][-1] * u.arcsec)
×
2248
            # check for partial spectra (for coronagraphs only)
2249
            if not (mode["syst"]["occulter"]):
1✔
2250
                IWA_max = mode["IWA"] * (1.0 + mode["BW"] / 2.0)
1✔
2251
                OWA_min = mode["OWA"] * (1.0 - mode["BW"] / 2.0)
1✔
2252
                char[char] = (WAchar < IWA_max) | (WAchar > OWA_min)
1✔
2253
                characterized[char] = -1
1✔
2254
            # encode results in spectra lists (only for planets, not FA)
2255
            charplans = characterized[:-1] if FA else characterized
1✔
2256
            self.fullSpectra[pInds[charplans == 1]] += 1
1✔
2257
            self.partialSpectra[pInds[charplans == -1]] += 1
1✔
2258

2259
        if self.make_debug_bird_plots:
1✔
2260
            from tools.obs_plot import obs_plot
×
2261

2262
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2263

2264
        return characterized.astype(int), fZ, systemParams, SNR, intTime
1✔
2265

2266
    def calc_signal_noise(
1✔
2267
        self, sInd, pInds, t_int, mode, fZ=None, fEZ=None, dMag=None, WA=None
2268
    ):
2269
        """Calculates the signal and noise fluxes for a given time interval. Called
2270
        by observation_detection and observation_characterization methods in the
2271
        SurveySimulation module.
2272

2273
        Args:
2274
            sInd (int):
2275
                Integer index of the star of interest
2276
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2277
                Integration time interval in units of day
2278
            pInds (int):
2279
                Integer indices of the planets of interest
2280
            mode (dict):
2281
                Selected observing mode (from OpticalSystem)
2282
            fZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2283
                Surface brightness of local zodiacal light in units of 1/arcsec2
2284
            fEZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2285
                Surface brightness of exo-zodiacal light in units of 1/arcsec2
2286
            dMag (~numpy.ndarray(float)):
2287
                Differences in magnitude between planets and their host star
2288
            WA (~astropy.units.Quantity(~numpy.ndarray(float))):
2289
                Working angles of the planets of interest in units of arcsec
2290

2291
        Returns:
2292
            tuple:
2293
                Signal (float):
2294
                    Counts of signal
2295
                Noise (float):
2296
                    Counts of background noise variance
2297

2298
        """
2299

2300
        OS = self.OpticalSystem
1✔
2301
        ZL = self.ZodiacalLight
1✔
2302
        TL = self.TargetList
1✔
2303
        SU = self.SimulatedUniverse
1✔
2304
        Obs = self.Observatory
1✔
2305
        TK = self.TimeKeeping
1✔
2306

2307
        # calculate optional parameters if not provided
2308
        fZ = (
1✔
2309
            fZ
2310
            if (fZ is not None)
2311
            else ZL.fZ(Obs, TL, sInd, TK.currentTimeAbs.copy(), mode)
2312
        )
2313
        fEZ = fEZ if (fEZ is not None) else SU.fEZ[pInds]
1✔
2314

2315
        # if lucky_planets, use lucky planet params for dMag and WA
2316
        if SU.lucky_planets and mode in list(
1✔
2317
            filter(lambda mode: "spec" in mode["inst"]["name"], OS.observingModes)
2318
        ):
2319
            phi = (1 / np.pi) * np.ones(len(SU.d))
×
2320
            dMag = deltaMag(SU.p, SU.Rp, SU.d, phi)[pInds]  # delta magnitude
×
2321
            WA = np.arctan(SU.a / TL.dist[SU.plan2star]).to("arcsec")[
×
2322
                pInds
2323
            ]  # working angle
2324
        else:
2325
            dMag = dMag if (dMag is not None) else SU.dMag[pInds]
1✔
2326
            WA = WA if (WA is not None) else SU.WA[pInds]
1✔
2327

2328
        # initialize Signal and Noise arrays
2329
        Signal = np.zeros(len(pInds))
1✔
2330
        Noise = np.zeros(len(pInds))
1✔
2331

2332
        # find observable planets wrt IWA-OWA range
2333
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2334

2335
        if np.any(obs):
1✔
2336
            # find electron counts
2337
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(
1✔
2338
                TL, sInd, fZ, fEZ[obs], dMag[obs], WA[obs], mode, TK=TK
2339
            )
2340
            # calculate signal and noise levels (based on Nemati14 formula)
2341
            Signal[obs] = (C_p * t_int).decompose().value
1✔
2342
            Noise[obs] = np.sqrt((C_b * t_int + (C_sp * t_int) ** 2).decompose().value)
1✔
2343

2344
        return Signal, Noise
1✔
2345

2346
    def update_occulter_mass(self, DRM, sInd, t_int, skMode):
1✔
2347
        """Updates the occulter wet mass in the Observatory module, and stores all
2348
        the occulter related values in the DRM array.
2349

2350
        Args:
2351
            DRM (dict):
2352
                Design Reference Mission, contains the results of one complete
2353
                observation (detection and characterization)
2354
            sInd (int):
2355
                Integer index of the star of interest
2356
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2357
                Selected star integration time (for detection or characterization)
2358
                in units of day
2359
            skMode (str):
2360
                Station keeping observing mode type ('det' or 'char')
2361

2362
        Returns:
2363
            dict:
2364
                Design Reference Mission dictionary, contains the results of one
2365
                complete observation
2366

2367
        """
2368

2369
        TL = self.TargetList
×
2370
        Obs = self.Observatory
×
2371
        TK = self.TimeKeeping
×
2372

2373
        assert skMode in ("det", "char"), "Observing mode type must be 'det' or 'char'."
×
2374

2375
        # decrement mass for station-keeping
2376
        dF_lateral, dF_axial, intMdot, mass_used, deltaV = Obs.mass_dec_sk(
×
2377
            TL, sInd, TK.currentTimeAbs.copy(), t_int
2378
        )
2379

2380
        DRM[skMode + "_dV"] = deltaV.to("m/s")
×
2381
        DRM[skMode + "_mass_used"] = mass_used.to("kg")
×
2382
        DRM[skMode + "_dF_lateral"] = dF_lateral.to("N")
×
2383
        DRM[skMode + "_dF_axial"] = dF_axial.to("N")
×
2384
        # update spacecraft mass
2385
        Obs.scMass = Obs.scMass - mass_used
×
2386
        DRM["scMass"] = Obs.scMass.to("kg")
×
2387
        if Obs.twotanks:
×
2388
            Obs.skMass = Obs.skMass - mass_used
×
2389
            DRM["skMass"] = Obs.skMass.to("kg")
×
2390

2391
        return DRM
×
2392

2393
    def reset_sim(self, genNewPlanets=True, rewindPlanets=True, seed=None):
1✔
2394
        """Performs a full reset of the current simulation.
2395

2396
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2397
        objects with their own outspecs.
2398

2399
        Args:
2400
            genNewPlanets (bool):
2401
                Generate all new planets based on the original input specification.
2402
                If False, then the original planets will remain. Setting to True forces
2403
                ``rewindPlanets`` to be True as well. Defaults True.
2404
            rewindPlanets (bool):
2405
                Reset the current set of planets to their original orbital phases.
2406
                If both genNewPlanets and rewindPlanet  are False, the original planets
2407
                will be retained and will not be rewound to their initial starting
2408
                locations (i.e., all systems will remain at the times they were at the
2409
                end of the last run, thereby effectively randomizing planet phases.
2410
                Defaults True.
2411
            seed (int, optional):
2412
                Random seed to use for all random number generation. If None (default)
2413
                a new random seed will be generated when re-initializing the
2414
                SurveySimulation.
2415

2416
        """
2417

2418
        SU = self.SimulatedUniverse
1✔
2419
        TK = self.TimeKeeping
1✔
2420
        TL = self.TargetList
1✔
2421

2422
        # re-initialize SurveySimulation arrays
2423
        specs = self._outspec
1✔
2424
        specs["modules"] = self.modules
1✔
2425

2426
        if seed is None:  # pop the seed so a new one is generated
1✔
2427
            if "seed" in specs:
1✔
2428
                specs.pop("seed")
1✔
2429
        else:  # if seed is provided, replace seed with provided seed
2430
            specs["seed"] = seed
×
2431

2432
        # reset mission time, re-init surveysim and observatory
2433
        TK.__init__(**TK._outspec)
1✔
2434
        self.__init__(**specs)
1✔
2435
        self.Observatory.__init__(**self.Observatory._outspec)
1✔
2436

2437
        # generate new planets if requested (default)
2438
        if genNewPlanets:
1✔
2439
            TL.stellar_mass()
1✔
2440
            TL.I = TL.gen_inclinations(TL.PlanetPopulation.Irange)  # noqa: E741
1✔
2441
            SU.gen_physical_properties(**SU._outspec)
1✔
2442
            rewindPlanets = True
1✔
2443
        # re-initialize systems if requested (default)
2444
        if rewindPlanets:
1✔
2445
            SU.init_systems()
1✔
2446

2447
        # reset helper arrays
2448
        self.initializeStorageArrays()
1✔
2449

2450
        self.vprint("Simulation reset.")
1✔
2451

2452
    def genOutSpec(
1✔
2453
        self,
2454
        tofile: Optional[str] = None,
2455
        starting_outspec: Optional[Dict[str, Any]] = None,
2456
        modnames: bool = False,
2457
    ) -> Dict[str, Any]:
2458
        """Join all _outspec dicts from all modules into one output dict
2459
        and optionally write out to JSON file on disk.
2460

2461
        Args:
2462
            tofile (str, optional):
2463
                Name of the file containing all output specifications (outspecs).
2464
                Defaults to None.
2465
            starting_outspec (dict, optional):
2466
                Initial outspec (from MissionSim). Defaults to None.
2467
            modnames (bool):
2468
                If True, populate outspec dictionary with the module it originated from,
2469
                instead of the actual value of the keyword. Defaults False.
2470

2471
        Returns:
2472
            dict:
2473
                Dictionary containing the full :ref:`sec:inputspec`, including all
2474
                filled-in default values. Combination of all individual module _outspec
2475
                attributes.
2476

2477
        """
2478

2479
        # start with our own outspec
2480
        if modnames:
1✔
2481
            out = copy.copy(self._outspec)
1✔
2482
            for k in out:
1✔
2483
                out[k] = self.__class__.__name__
1✔
2484
        else:
2485
            out = copy.deepcopy(self._outspec)
1✔
2486

2487
        # Add any provided other outspec
2488
        if starting_outspec is not None:
1✔
2489
            out.update(starting_outspec)
1✔
2490

2491
        # add in all modules _outspec's
2492
        for module in self.modules.values():
1✔
2493
            if modnames:
1✔
2494
                tmp = copy.copy(module._outspec)
1✔
2495
                for k in tmp:
1✔
2496
                    tmp[k] = module.__class__.__name__
1✔
2497
            else:
2498
                tmp = module._outspec
1✔
2499
            out.update(tmp)
1✔
2500

2501
        # add in the specific module names used
2502
        out["modules"] = {}
1✔
2503
        for mod_name, module in self.modules.items():
1✔
2504
            # find the module file
2505
            mod_name_full = module.__module__
1✔
2506
            if mod_name_full.startswith("EXOSIMS"):
1✔
2507
                # take just its short name if it is in EXOSIMS
2508
                mod_name_short = mod_name_full.split(".")[-1]
1✔
2509
            else:
2510
                # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2511
                mod_name_short = re.sub(
×
2512
                    r"\.pyc$", ".py", inspect.getfile(module.__class__)
2513
                )
2514
            out["modules"][mod_name] = mod_name_short
1✔
2515
        # add catalog name
2516
        if self.TargetList.keepStarCatalog:
1✔
2517
            module = self.TargetList.StarCatalog
×
2518
            mod_name_full = module.__module__
×
2519
            if mod_name_full.startswith("EXOSIMS"):
×
2520
                # take just its short name if it is in EXOSIMS
2521
                mod_name_short = mod_name_full.split(".")[-1]
×
2522
            else:
2523
                # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2524
                mod_name_short = re.sub(
×
2525
                    r"\.pyc$", ".py", inspect.getfile(module.__class__)
2526
                )
2527
            out["modules"][mod_name] = mod_name_short
×
2528
        else:
2529
            out["modules"][
1✔
2530
                "StarCatalog"
2531
            ] = self.TargetList.StarCatalog  # we just copy the StarCatalog string
2532

2533
        # if we don't know about the SurveyEnsemble, just write a blank to the output
2534
        if "SurveyEnsemble" not in out["modules"]:
1✔
2535
            out["modules"]["SurveyEnsemble"] = " "
1✔
2536

2537
        # add in the SVN/Git revision
2538
        path = os.path.split(inspect.getfile(self.__class__))[0]
1✔
2539
        path = os.path.split(os.path.split(path)[0])[0]
1✔
2540
        # handle case where EXOSIMS was imported from the working directory
2541
        if path == "":
1✔
2542
            path = os.getcwd()
×
2543
        # comm = "git -C " + path + " log -1"
2544
        comm = "git --git-dir=%s --work-tree=%s log -1" % (
1✔
2545
            os.path.join(path, ".git"),
2546
            path,
2547
        )
2548
        rev = subprocess.Popen(
1✔
2549
            comm, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
2550
        )
2551
        (gitRev, err) = rev.communicate()
1✔
2552
        gitRev = gitRev.decode("utf-8")
1✔
2553
        if isinstance(gitRev, str) & (len(gitRev) > 0):
1✔
2554
            tmp = re.compile(
1✔
2555
                r"\S*(commit [0-9a-fA-F]+)\n[\s\S]*Date: ([\S ]*)\n"
2556
            ).match(gitRev)
2557
            if tmp:
1✔
2558
                out["Revision"] = "Github " + tmp.groups()[0] + " " + tmp.groups()[1]
1✔
2559
        else:
2560
            rev = subprocess.Popen(
×
2561
                "svn info " + path + "| grep \"Revision\" | awk '{print $2}'",
2562
                stdout=subprocess.PIPE,
2563
                stderr=subprocess.PIPE,
2564
                shell=True,
2565
            )
2566
            (svnRev, err) = rev.communicate()
×
2567
            if isinstance(svnRev, str) & (len(svnRev) > 0):
×
2568
                out["Revision"] = "SVN revision is " + svnRev[:-1]
×
2569
            else:
2570
                out["Revision"] = "Not a valid Github or SVN revision."
×
2571

2572
        # dump to file
2573
        if tofile is not None:
1✔
2574
            with open(tofile, "w") as outfile:
1✔
2575
                json.dump(
1✔
2576
                    out,
2577
                    outfile,
2578
                    sort_keys=True,
2579
                    indent=4,
2580
                    ensure_ascii=False,
2581
                    separators=(",", ": "),
2582
                    default=array_encoder,
2583
                )
2584

2585
        return out
1✔
2586

2587
    def generateHashfName(self, specs):
1✔
2588
        """Generate cached file Hashname
2589

2590
        Requires a .XXX appended to end of hashname for each individual use case
2591

2592
        Args:
2593
            specs (dict):
2594
                :ref:`sec:inputspec`
2595

2596
        Returns:
2597
            str:
2598
                Unique indentifier string for cahce products from this set of modules
2599
                and inputs
2600
        """
2601
        # Allows cachefname to be predefined
2602
        if "cachefname" in specs:
1✔
2603
            return specs["cachefname"]
×
2604

2605
        cachefname = ""  # declares cachefname
1✔
2606
        mods = ["Completeness", "TargetList", "OpticalSystem"]  # modules to look at
1✔
2607
        tmp = (
1✔
2608
            self.Completeness.PlanetPopulation.__class__.__name__
2609
            + self.Completeness.PlanetPhysicalModel.__class__.__name__
2610
            + self.PlanetPopulation.__class__.__name__
2611
            + self.SimulatedUniverse.__class__.__name__
2612
            + self.PlanetPhysicalModel.__class__.__name__
2613
        )
2614

2615
        if "selectionMetric" in specs:
1✔
2616
            tmp += specs["selectionMetric"]
×
2617
        if "Izod" in specs:
1✔
2618
            tmp += specs["Izod"]
×
2619
        if "maxiter" in specs:
1✔
2620
            tmp += str(specs["maxiter"])
×
2621
        if "ftol" in specs:
1✔
2622
            tmp += str(specs["ftol"])
×
2623
        if "missionLife" in specs:
1✔
2624
            tmp += str(specs["missionLife"])
1✔
2625
        if "missionPortion" in specs:
1✔
2626
            tmp += str(specs["missionPortion"])
1✔
2627
        if "smaknee" in specs:
1✔
2628
            tmp += str(specs["smaknee"])
×
2629
        if "koAngleMax" in specs:
1✔
2630
            tmp += str(specs["koAngleMax"])
×
2631
        tmp += str(np.sum(self.Completeness.PlanetPopulation.arange.value))
1✔
2632
        tmp += str(np.sum(self.Completeness.PlanetPopulation.Rprange.value))
1✔
2633
        tmp += str(np.sum(self.Completeness.PlanetPopulation.erange))
1✔
2634
        tmp += str(
1✔
2635
            self.Completeness.PlanetPopulation.PlanetPhysicalModel.whichPlanetPhaseFunction  # noqa: E501
2636
        )
2637
        tmp += str(np.sum(self.PlanetPopulation.arange.value))
1✔
2638
        tmp += str(np.sum(self.PlanetPopulation.Rprange.value))
1✔
2639
        tmp += str(np.sum(self.PlanetPopulation.erange))
1✔
2640
        tmp += str(self.PlanetPopulation.PlanetPhysicalModel.whichPlanetPhaseFunction)
1✔
2641

2642
        for mod in mods:
1✔
2643
            cachefname += self.modules[mod].__module__.split(".")[
1✔
2644
                -1
2645
            ]  # add module name to end of cachefname
2646
        cachefname += hashlib.md5(
1✔
2647
            (str(self.TargetList.Name) + tmp).encode("utf-8")
2648
        ).hexdigest()  # turn cachefname into hashlib
2649
        cachefname = os.path.join(
1✔
2650
            self.cachedir, cachefname + os.extsep
2651
        )  # join into filepath and fname
2652
        # Needs file terminator (.starkt0, .t0, etc) appended done by each individual
2653
        # use case.
2654
        return cachefname
1✔
2655

2656
    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
1✔
2657
        """Helper method for Overloading Revisit Filtering
2658

2659
        Args:
2660
            sInds (~numpy.ndarray(int)):
2661
                Indices of stars still in observation list
2662
            tmpCurrentTimeNorm (astropy.units.Quantity):
2663
                Normalized simulation time
2664

2665
        Returns:
2666
            ~numpy.ndarray(int):
2667
                indices of stars still in observation list
2668
        """
2669

2670
        tovisit = np.zeros(self.TargetList.nStars, dtype=bool)
1✔
2671
        if len(sInds) > 0:
1✔
2672
            # Check that no star has exceeded the number of revisits and the indicies
2673
            # of all considered stars have minimum number of observations
2674
            # This should prevent revisits so long as all stars have not
2675
            # been observed
2676
            tovisit[sInds] = (self.starVisits[sInds] == min(self.starVisits[sInds])) & (
1✔
2677
                self.starVisits[sInds] < self.nVisitsMax
2678
            )
2679
            if self.starRevisit.size != 0:
1✔
2680
                ind_rev = self.revisit_inds(sInds, tmpCurrentTimeNorm)
1✔
2681
                tovisit[ind_rev] = self.starVisits[ind_rev] < self.nVisitsMax
1✔
2682
            sInds = np.where(tovisit)[0]
1✔
2683
        return sInds
1✔
2684

2685
    def is_earthlike(self, plan_inds, sInd):
1✔
2686
        """Is the planet earthlike? Returns boolean array that's True for Earth-like
2687
        planets.
2688

2689

2690
        Args:
2691
            plan_inds(~numpy.ndarray(int)):
2692
                Planet indices
2693
            sInd (int):
2694
                Star index
2695

2696
        Returns:
2697
            ~numpy.ndarray(bool):
2698
                Array of same dimension as plan_inds input that's True for Earthlike
2699
                planets and false otherwise.
2700
        """
2701
        TL = self.TargetList
1✔
2702
        SU = self.SimulatedUniverse
1✔
2703
        PPop = self.PlanetPopulation
1✔
2704

2705
        # extract planet and star properties
2706
        Rp_plan = SU.Rp[plan_inds].value
1✔
2707
        L_star = TL.L[sInd]
1✔
2708
        if PPop.scaleOrbits:
1✔
2709
            a_plan = (SU.a[plan_inds] / np.sqrt(L_star)).value
×
2710
        else:
2711
            a_plan = (SU.a[plan_inds]).value
1✔
2712
        # Definition: planet radius (in earth radii) and solar-equivalent luminosity
2713
        # must be between the given bounds.
2714
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
1✔
2715
        # We use the numpy versions so that plan_ind can be a numpy vector.
2716
        return np.logical_and(
1✔
2717
            np.logical_and(Rp_plan >= Rp_plan_lo, Rp_plan <= 1.4),
2718
            np.logical_and(a_plan >= 0.95, a_plan <= 1.67),
2719
        )
2720

2721
    def find_known_plans(self):
1✔
2722
        """
2723
        Find and return list of known RV stars and list of stars with earthlike planets
2724
        based on info from David Plavchan dated 12/24/2018
2725
        """
2726
        TL = self.TargetList
×
2727
        SU = self.SimulatedUniverse
×
2728
        PPop = self.PlanetPopulation
×
2729
        L_star = TL.L[SU.plan2star]
×
2730

2731
        c = 28.4 * u.m / u.s
×
2732
        Mj = 317.8 * u.earthMass
×
2733
        Mpj = SU.Mp / Mj  # planet masses in jupiter mass units
×
2734
        Ms = TL.MsTrue[SU.plan2star]
×
2735
        Teff = TL.Teff[SU.plan2star]
×
2736
        mu = const.G * (SU.Mp + Ms)
×
2737
        T = (2.0 * np.pi * np.sqrt(SU.a**3 / mu)).to(u.yr)
×
2738
        e = SU.e
×
2739

2740
        # pinds in correct temp range
2741
        t_filt = np.where((Teff.value > 3000) & (Teff.value < 6800))[0]
×
2742

2743
        K = (
×
2744
            (c / np.sqrt(1 - e[t_filt]))
2745
            * Mpj[t_filt]
2746
            * np.sin(SU.I[t_filt])
2747
            * Ms[t_filt] ** (-2 / 3)
2748
            * T[t_filt] ** (-1 / 3)
2749
        )
2750

2751
        K_filter = (T[t_filt].to(u.d) / 10**4).value  # create period-filter
×
2752
        # if period-filter value is lower than .03, set to .03
2753
        K_filter[np.where(K_filter < 0.03)[0]] = 0.03
×
2754
        k_filt = t_filt[np.where(K.value > K_filter)[0]]  # pinds in the correct K range
×
2755

2756
        if PPop.scaleOrbits:
×
2757
            a_plan = (SU.a / np.sqrt(L_star)).value
×
2758
        else:
2759
            a_plan = SU.a.value
×
2760

2761
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
×
2762

2763
        # pinds in habitable zone
2764
        a_filt = k_filt[np.where((a_plan[k_filt] > 0.95) & (a_plan[k_filt] < 1.67))[0]]
×
2765
        # rocky planets
2766
        r_filt = a_filt[
×
2767
            np.where(
2768
                (SU.Rp.value[a_filt] >= Rp_plan_lo[a_filt])
2769
                & (SU.Rp.value[a_filt] < 1.4)
2770
            )[0]
2771
        ]
2772
        self.known_earths = np.union1d(self.known_earths, r_filt).astype(int)
×
2773

2774
        # these are known_rv stars with earths around them
2775
        known_stars = np.unique(SU.plan2star[k_filt])  # these are known_rv stars
×
2776
        known_rocky = np.unique(SU.plan2star[r_filt])
×
2777

2778
        # if include_known_RV, then filter out all other sInds
2779
        if self.include_known_RV is not None:
×
2780
            HIP_sInds = np.where(np.in1d(TL.Name, self.include_known_RV))[0]
×
2781
            known_stars = np.intersect1d(HIP_sInds, known_stars)
×
2782
            known_rocky = np.intersect1d(HIP_sInds, known_rocky)
×
2783
        return known_stars.astype(int), known_rocky.astype(int)
×
2784

2785
    def find_char_SNR(self, sInd, pIndsChar, startTime, intTime, mode):
1✔
2786
        """Finds the SNR achieved by an observing mode after intTime days
2787

2788
        The observation window (which includes settling and overhead times)
2789
        is a superset of the integration window (in which photons are collected).
2790

2791
        The observation window begins at startTime. The integration window
2792
        begins at startTime + Obs.settlingTime + mode["ohTime"],
2793
        and the integration itself has a duration of intTime.
2794

2795
        Args:
2796
            sInd (int):
2797
                Integer index of the star of interest
2798
            pIndsChar (int numpy.ndarray):
2799
                Observable planets to characterize. Place (-1) at end to put
2800
                False Alarm parameters at end of returned tuples.
2801
            startTime (astropy.units.Quantity):
2802
                Beginning of observation window in units of day.
2803
            intTime (astropy.units.Quantity):
2804
                Selected star characterization integration time in units of day.
2805
            mode (dict):
2806
                Observing mode for the characterization
2807

2808
        Returns:
2809
            tuple:
2810
                SNR (float numpy.ndarray):
2811
                    Characterization signal-to-noise ratio of the observable planets.
2812
                    Defaults to None. [TBD]
2813
                fZ (astropy.units.Quantity):
2814
                    Surface brightness of local zodiacal light in units of 1/arcsec2
2815
                systemParams (dict):
2816
                    Dictionary of time-dependent planet properties averaged over the
2817
                    duration of the integration.
2818

2819
        """
2820

2821
        OS = self.OpticalSystem
×
2822
        ZL = self.ZodiacalLight
×
2823
        TL = self.TargetList
×
2824
        SU = self.SimulatedUniverse
×
2825
        Obs = self.Observatory
×
2826
        TK = self.TimeKeeping
×
2827

2828
        # time at start of integration window
2829
        currentTimeNorm = startTime
×
2830
        currentTimeAbs = TK.missionStart + startTime
×
2831

2832
        # first, calculate SNR for observable planets (without false alarm)
2833
        planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
×
2834
        SNRplans = np.zeros(len(planinds))
×
2835
        if len(planinds) > 0:
×
2836
            # initialize arrays for SNR integration
2837
            fZs = np.zeros(self.ntFlux) / u.arcsec**2.0
×
2838
            systemParamss = np.empty(self.ntFlux, dtype="object")
×
2839
            Ss = np.zeros((self.ntFlux, len(planinds)))
×
2840
            Ns = np.zeros((self.ntFlux, len(planinds)))
×
2841
            # integrate the signal (planet flux) and noise
2842
            dt = intTime / float(self.ntFlux)
×
2843
            timePlus = (
×
2844
                Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2845
            )  # accounts for the time since the current time
2846
            for i in range(self.ntFlux):
×
2847
                # calculate signal and noise (electron count rates)
2848
                if SU.lucky_planets:
×
2849
                    fZs[i] = ZL.fZ(Obs, TL, sInd, currentTimeAbs, mode)[0]
×
2850
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
×
2851
                        sInd, planinds, dt, mode, fZ=fZs[i]
2852
                    )
2853
                # allocate first half of dt
2854
                timePlus += dt / 2.0
×
2855
                # calculate current zodiacal light brightness
2856
                fZs[i] = ZL.fZ(Obs, TL, sInd, currentTimeAbs + timePlus, mode)[0]
×
2857
                # propagate the system to match up with current time
2858
                SU.propag_system(
×
2859
                    sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2860
                )
2861
                self.propagTimes[sInd] = currentTimeNorm + timePlus
×
2862
                # time-varying planet params (WA, dMag, phi, fEZ, d)
2863
                systemParamss[i] = SU.dump_system_params(sInd)
×
2864
                # calculate signal and noise (electron count rates)
2865
                if not SU.lucky_planets:
×
2866
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
×
2867
                        sInd, planinds, dt, mode, fZ=fZs[i]
2868
                    )
2869
                # allocate second half of dt
2870
                timePlus += dt / 2.0
×
2871

2872
            # average output parameters
2873
            fZ = np.mean(fZs)
×
2874
            systemParams = {
×
2875
                key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2876
                / float(self.ntFlux)
2877
                for key in sorted(systemParamss[0])
2878
            }
2879
            # calculate planets SNR
2880
            S = Ss.sum(0)
×
2881
            N = Ns.sum(0)
×
2882
            SNRplans[N > 0] = S[N > 0] / N[N > 0]
×
2883
            # allocate extra time for timeMultiplier
2884

2885
        # if only a FA, just save zodiacal brightness
2886
        # in the middle of the integration
2887
        else:
2888
            totTime = intTime * (mode["timeMultiplier"])
×
2889
            fZ = ZL.fZ(Obs, TL, sInd, currentTimeAbs.copy() + totTime / 2.0, mode)[0]
×
2890

2891
        # calculate the false alarm SNR (if any)
2892
        SNRfa = []
×
2893
        if pIndsChar[-1] == -1:
×
2894
            # Note: these attributes may not exist for all schedulers
2895
            fEZ = self.lastDetected[sInd, 1][-1] / u.arcsec**2.0
×
2896
            dMag = self.lastDetected[sInd, 2][-1]
×
2897
            WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
2898
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, fEZ, dMag, WA, mode)
×
2899
            S = (C_p * intTime).decompose().value
×
2900
            N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2901
            SNRfa = S / N if N > 0.0 else 0.0
×
2902

2903
        # SNR vector is of length (#planets), plus 1 if FA
2904
        # This routine leaves SNR = 0 if unknown or not found
2905
        pInds = np.where(SU.plan2star == sInd)[0]
×
2906
        FA_present = pIndsChar[-1] == -1
×
2907
        # there will be one SNR for each entry of pInds_FA
2908
        pInds_FA = np.append(pInds, [-1] if FA_present else np.zeros(0, dtype=int))
×
2909

2910
        # boolean index vector into SNR
2911
        #   True iff we have computed a SNR for that planet
2912
        #   False iff we didn't find an SNR (and will leave 0 there)
2913
        #   if FA_present, SNR_plug_in[-1] = True
2914
        SNR_plug_in = np.isin(pInds_FA, pIndsChar)
×
2915

2916
        # save all SNRs (planets and FA) to one array
2917
        SNR = np.zeros(len(pInds_FA))
×
2918
        # plug in the SNR's we computed (pIndsChar and optional FA)
2919
        SNR[SNR_plug_in] = np.append(SNRplans, SNRfa)
×
2920

2921
        return SNR, fZ, systemParams
×
2922

2923
    def revisit_inds(self, sInds, tmpCurrentTimeNorm):
1✔
2924
        """Return subset of star indices that are scheduled for revisit.
2925

2926
        Args:
2927
            sInds (~numpy.ndarray(int)):
2928
                Indices of stars to consider
2929
            tmpCurrentTimeNorm (astropy.units.Quantity):
2930
                Normalized simulation time
2931

2932
        Returns:
2933
            ~numpy.ndarray(int):
2934
                Indices of stars whose revisit is scheduled for within `self.dt_max` of
2935
                the current time
2936

2937
        """
2938
        dt_rev = np.abs(self.starRevisit[:, 1] * u.day - tmpCurrentTimeNorm)
1✔
2939
        ind_rev = [
1✔
2940
            int(x) for x in self.starRevisit[dt_rev < self.dt_max, 0] if x in sInds
2941
        ]
2942

2943
        return ind_rev
1✔
2944

2945
    def keepout_filter(self, sInds, startTimes, endTimes, koMap):
1✔
2946
        """Filters stars not observable due to keepout
2947

2948
        Args:
2949
            sInds (~numpy.ndarray(int)):
2950
                indices of stars still in observation list
2951
            startTimes (~numpy.ndarray(float)):
2952
                start times of observations
2953
            endTimes (~numpy.ndarray(float)):
2954
                end times of observations
2955
            koMap (~numpy.ndarray(bool)):
2956
                map where True is for a target unobstructed and observable,
2957
                False is for a target obstructed and unobservable due to keepout zone
2958

2959
        Returns:
2960
            ~numpy.ndarray(int):
2961
                sInds - filtered indices of stars still in observation list
2962

2963
        """
2964
        # find the indices of keepout times that pertain to the start and end times of
2965
        # observation
2966
        startInds = np.searchsorted(self.koTimes.value, startTimes)
×
2967
        endInds = np.searchsorted(self.koTimes.value, endTimes)
×
2968

2969
        # boolean array of available targets (unobstructed between observation start and
2970
        # end time) we include a -1 in the start and +1 in the end to include keepout
2971
        # days enclosing start and end times
2972
        availableTargets = np.array(
×
2973
            [
2974
                np.all(
2975
                    koMap[
2976
                        sInd,
2977
                        max(startInds[sInd] - 1, 0) : min(
2978
                            endInds[sInd] + 1, len(self.koTimes.value) + 1
2979
                        ),
2980
                    ]
2981
                )
2982
                for sInd in sInds
2983
            ],
2984
            dtype="bool",
2985
        )
2986

2987
        return sInds[availableTargets]
×
2988

2989

2990
def array_encoder(obj):
1✔
2991
    r"""Encodes numpy arrays, astropy Times, and astropy Quantities, into JSON.
2992

2993
    Called from json.dump for types that it does not already know how to represent,
2994
    like astropy Quantity's, numpy arrays, etc.  The json.dump() method encodes types
2995
    like ints, strings, and lists itself, so this code does not see these types.
2996
    Likewise, this routine can and does return such objects, which is OK as long as
2997
    they unpack recursively into types for which encoding is known
2998

2999
    Args:
3000
        obj (Any):
3001
            Object to encode.
3002

3003
        Returns:
3004
            Any:
3005
                Encoded object
3006

3007
    """
3008

3009
    from astropy.coordinates import SkyCoord
1✔
3010
    from astropy.time import Time
1✔
3011

3012
    if isinstance(obj, Time):
1✔
3013
        # astropy Time -> time string
3014
        return obj.fits  # isot also makes sense here
×
3015
    if isinstance(obj, u.quantity.Quantity):
1✔
3016
        # note: it is possible to have a numpy ndarray wrapped in a Quantity.
3017
        # NB: alternatively, can return (obj.value, obj.unit.name)
3018
        return obj.value
×
3019
    if isinstance(obj, SkyCoord):
1✔
3020
        return dict(
×
3021
            lon=obj.heliocentrictrueecliptic.lon.value,
3022
            lat=obj.heliocentrictrueecliptic.lat.value,
3023
            distance=obj.heliocentrictrueecliptic.distance.value,
3024
        )
3025
    if isinstance(obj, (np.ndarray, np.number)):
1✔
3026
        # ndarray -> list of numbers
3027
        return obj.tolist()
1✔
3028
    if isinstance(obj, complex):
1✔
3029
        # complex -> (real, imag) pair
3030
        return [obj.real, obj.imag]
×
3031
    if callable(obj):
1✔
3032
        # this case occurs for interpolants like PSF and QE
3033
        # We cannot simply "write" the function to JSON, so we make up a string
3034
        # to keep from throwing an error.
3035
        # The fix is simple: when generating the interpolant, add a _outspec attribute
3036
        # to the function (or the lambda), containing (e.g.) the fits filename, or the
3037
        # explicit number -- whatever string was used.  Then, here, check for that
3038
        # attribute and write it out instead of this dummy string.  (Attributes can
3039
        # be transparently attached to python functions, even lambda's.)
3040
        return "interpolant_function"
1✔
3041
    if isinstance(obj, set):
×
3042
        return list(obj)
×
3043
    if isinstance(obj, bytes):
×
3044
        return obj.decode()
×
3045
    # an EXOSIMS object
3046
    if hasattr(obj, "_modtype"):
×
3047
        return obj.__dict__
×
3048
    # an object for which no encoding is defined yet
3049
    #   as noted above, ordinary types (lists, ints, floats) do not take this path
3050
    raise ValueError("Could not JSON-encode an object of type %s" % type(obj))
×
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc