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

dsavransky / EXOSIMS / 14536111075

18 Apr 2025 01:44PM UTC coverage: 65.708% (+0.2%) from 65.466%
14536111075

push

github

web-flow
Merge pull request #402 from dsavransky/exozodi_fix

Exozodi fix. Extensive re-plumbing of all things related to exozodi.

210 of 269 new or added lines in 26 files covered. (78.07%)

14 existing lines in 7 files now uncovered.

9665 of 14709 relevant lines covered (65.71%)

0.66 hits per line

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

62.63
/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 sys
1✔
11
import time
1✔
12
from pathlib import Path
1✔
13
from typing import Any, Dict, Optional
1✔
14

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

20
from EXOSIMS.util.deltaMag import deltaMag
1✔
21
from EXOSIMS.util.get_dirs import get_cache_dir
1✔
22
from EXOSIMS.util.get_module import get_module
1✔
23
from EXOSIMS.util.vprint import vprint
1✔
24
from EXOSIMS.util._numpy_compat import copy_if_needed
1✔
25
from EXOSIMS.util.version_util import get_version
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
        # start the outspec
213
        self._outspec = {}
1✔
214

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

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

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

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

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

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

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

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

257
        # if any of the modules is a string, assume that they are all strings
258
        # and we need to initalize
259
        if isinstance(next(iter(specs["modules"].values())), str):
1✔
260
            # import desired module names (prototype or specific)
261
            self.SimulatedUniverse = get_module(
1✔
262
                specs["modules"]["SimulatedUniverse"], "SimulatedUniverse"
263
            )(**specs)
264
            self.Observatory = get_module(
1✔
265
                specs["modules"]["Observatory"], "Observatory"
266
            )(**specs)
267
            self.TimeKeeping = get_module(
1✔
268
                specs["modules"]["TimeKeeping"], "TimeKeeping"
269
            )(**specs)
270

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

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

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

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

311
                setattr(self, modName, specs["modules"][modName])
1✔
312

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

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

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

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

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

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

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

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

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

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

394
        # initialize arrays updated in run_sim()
395
        self.initializeStorageArrays()
1✔
396

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

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

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

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

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

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

432
        self._outspec["nofZ"] = nofZ
1✔
433

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

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

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

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

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

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

511
    def initializeStorageArrays(self):
1✔
512
        """
513
        Initialize all storage arrays based on # of stars and targets
514
        """
515

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

528
    def __str__(self):
1✔
529
        """String representation of the Survey Simulation object
530

531
        When the command 'print' is used on the Survey Simulation object, this
532
        method will return the values contained in the object
533

534
        """
535

536
        for att in self.__dict__:
1✔
537
            print("%s: %r" % (att, getattr(self, att)))
1✔
538

539
        return "Survey Simulation class object attributes"
1✔
540

541
    def run_sim(self):
1✔
542
        """Performs the survey simulation"""
543

544
        OS = self.OpticalSystem
1✔
545
        TL = self.TargetList
1✔
546
        SU = self.SimulatedUniverse
1✔
547
        Obs = self.Observatory
1✔
548
        TK = self.TimeKeeping
1✔
549

550
        # TODO: start using this self.currentSep
551
        # set occulter separation if haveOcculter
552
        if OS.haveOcculter:
1✔
553
            self.currentSep = Obs.occulterSep
×
554

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

568
        # begin Survey, and loop until mission is finished
569
        log_begin = "OB%s: survey beginning." % (TK.OBnumber)
1✔
570
        self.logger.info(log_begin)
1✔
571
        self.vprint(log_begin)
1✔
572
        t0 = time.time()
1✔
573
        sInd = None
1✔
574
        ObsNum = 0
1✔
575
        while not TK.mission_is_over(OS, Obs, det_mode):
1✔
576
            # acquire the NEXT TARGET star index and create DRM
577
            old_sInd = sInd  # used to save sInd if returned sInd is None
1✔
578
            DRM, sInd, det_intTime, waitTime = self.next_target(sInd, det_mode)
1✔
579

580
            if sInd is not None:
1✔
581
                ObsNum += (
1✔
582
                    1  # we're making an observation so increment observation number
583
                )
584

585
                if OS.haveOcculter:
1✔
586
                    # advance to start of observation (add slew time for selected target
587
                    _ = TK.advanceToAbsTime(TK.currentTimeAbs.copy() + waitTime)
×
588

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

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

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

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

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

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

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

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

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

805
        Returns:
806
            None
807
        """
808

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1036
        .. note::
1037

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

1040
        """
1041

1042
        TL = self.TargetList
1✔
1043

1044
        # assumed values for detection
1045
        fZ = self.ZodiacalLight.fZ(
1✔
1046
            self.Observatory, self.TargetList, sInds, startTimes, mode
1047
        )
1048
        # Use the default exozodi intensities
1049
        JEZ = TL.JEZ0[mode["hex"]][sInds]
1✔
1050
        dMag = TL.int_dMag[sInds]
1✔
1051
        WA = TL.int_WA[sInds]
1✔
1052

1053
        # save out file containing photon count info
1054
        if self.record_counts_path is not None and len(self.count_lines) == 0:
1✔
1055
            C_p, C_b, C_sp, C_extra = self.OpticalSystem.Cp_Cb_Csp(
×
1056
                self.TargetList, sInds, fZ, JEZ, dMag, WA, mode, returnExtra=True
1057
            )
1058
            import csv
×
1059

1060
            count_fpath = os.path.join(self.record_counts_path, "counts")
×
1061

1062
            if not os.path.exists(count_fpath):
×
1063
                os.mkdir(count_fpath)
×
1064

1065
            outfile = os.path.join(count_fpath, str(self.seed) + ".csv")
×
1066
            self.count_lines.append(
×
1067
                [
1068
                    "sInd",
1069
                    "HIPs",
1070
                    "C_F0",
1071
                    "C_p0",
1072
                    "C_sr",
1073
                    "C_z",
1074
                    "C_ez",
1075
                    "C_dc",
1076
                    "C_cc",
1077
                    "C_rn",
1078
                    "C_p",
1079
                    "C_b",
1080
                    "C_sp",
1081
                ]
1082
            )
1083

1084
            for i, sInd in enumerate(sInds):
×
1085
                self.count_lines.append(
×
1086
                    [
1087
                        sInd,
1088
                        self.TargetList.Name[sInd],
1089
                        C_extra["C_F0"][0].value,
1090
                        C_extra["C_sr"][i].value,
1091
                        C_extra["C_z"][i].value,
1092
                        C_extra["C_ez"][i].value,
1093
                        C_extra["C_dc"][i].value,
1094
                        C_extra["C_cc"][i].value,
1095
                        C_extra["C_rn"][i].value,
1096
                        C_p[i].value,
1097
                        C_b[i].value,
1098
                        C_sp[i].value,
1099
                    ]
1100
                )
1101

1102
            with open(outfile, "w") as csvfile:
×
1103
                c = csv.writer(csvfile)
×
1104
                c.writerows(self.count_lines)
×
1105

1106
        intTimes = self.OpticalSystem.calc_intTime(
1✔
1107
            self.TargetList, sInds, fZ, JEZ, dMag, WA, mode
1108
        )
1109
        intTimes[~np.isfinite(intTimes)] = 0 * u.d
1✔
1110

1111
        return intTimes
1✔
1112

1113
    def choose_next_target(self, old_sInd, sInds, slewTimes, intTimes):
1✔
1114
        """Helper method for method next_target to simplify alternative implementations.
1115

1116
        Given a subset of targets (pre-filtered by method next_target or some
1117
        other means), select the best next one. The prototype uses completeness
1118
        as the sole heuristic.
1119

1120
        Args:
1121
            old_sInd (int):
1122
                Index of the previous target star
1123
            sInds (~numpy.ndarray(int)):
1124
                Indices of available targets
1125
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1126
                slew times to all stars (must be indexed by sInds)
1127
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1128
                Integration times for detection in units of day
1129

1130
        Returns:
1131
            tuple:
1132
                sInd (int):
1133
                    Index of next target star
1134
                waitTime (:py:class:`~astropy.units.Quantity`):
1135
                    Some strategic amount of time to wait in case an occulter slew is
1136
                    desired (default is None)
1137

1138
        """
1139

1140
        Comp = self.Completeness
1✔
1141
        TL = self.TargetList
1✔
1142
        TK = self.TimeKeeping
1✔
1143
        OS = self.OpticalSystem
1✔
1144
        Obs = self.Observatory
1✔
1145
        allModes = OS.observingModes
1✔
1146

1147
        # cast sInds to array
1148
        sInds = np.array(sInds, ndmin=1, copy=copy_if_needed)
1✔
1149
        # calculate dt since previous observation
1150
        dt = TK.currentTimeNorm.copy() + slewTimes[sInds] - self.lastObsTimes[sInds]
1✔
1151
        # get dynamic completeness values
1152
        comps = Comp.completeness_update(TL, sInds, self.starVisits[sInds], dt)
1✔
1153
        # choose target with maximum completeness
1154
        sInd = np.random.choice(sInds[comps == max(comps)])
1✔
1155

1156
        # Check if exoplanetObsTime would be exceeded
1157
        mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
1158
        (
1✔
1159
            maxIntTimeOBendTime,
1160
            maxIntTimeExoplanetObsTime,
1161
            maxIntTimeMissionLife,
1162
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
1163
        maxIntTime = min(
1✔
1164
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1165
        )  # Maximum intTime allowed
1166
        intTimes2 = self.calc_targ_intTime(
1✔
1167
            np.array([sInd]), TK.currentTimeAbs.copy(), mode
1168
        )
1169
        if (
1✔
1170
            intTimes2 > maxIntTime
1171
        ):  # check if max allowed integration time would be exceeded
1172
            self.vprint("max allowed integration time would be exceeded")
×
1173
            sInd = None
×
1174
            waitTime = 1.0 * u.d
×
1175
            return sInd, waitTime
×
1176

1177
        return (
1✔
1178
            sInd,
1179
            slewTimes[sInd],
1180
        )  # if coronagraph or first sInd, waitTime will be 0 days
1181

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

1185
        Refines the selection of occulter slew times by filtering based on mission time
1186
        constraints and selecting the best slew time for each star. This method calls on
1187
        other occulter methods within SurveySimulation depending on how slew times were
1188
        calculated prior to calling this function (i.e. depending on which
1189
        implementation of the Observatory module is used).
1190

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

1209
        Returns:
1210
            tuple:
1211
                sInds (int):
1212
                    Indeces of next target star
1213
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1214
                    slew times to all stars (must be indexed by sInds)
1215
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1216
                    Integration times for detection in units of day
1217
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1218
                    Delta-V used to transfer to new star line of sight in unis of m/s
1219
        """
1220

1221
        Obs = self.Observatory
×
1222
        TL = self.TargetList
×
1223

1224
        # initializing arrays
1225
        obsTimeArray = np.zeros([TL.nStars, 50]) * u.d
×
1226
        intTimeArray = np.zeros([TL.nStars, 2]) * u.d
×
1227

1228
        for n in sInds:
×
1229
            obsTimeArray[n, :] = (
×
1230
                np.linspace(obsTimes[0, n].value, obsTimes[1, n].value, 50) * u.d
1231
            )
1232
        intTimeArray[sInds, 0] = self.calc_targ_intTime(
×
1233
            sInds, Time(obsTimeArray[sInds, 0], format="mjd", scale="tai"), mode
1234
        )
1235
        intTimeArray[sInds, 1] = self.calc_targ_intTime(
×
1236
            sInds, Time(obsTimeArray[sInds, -1], format="mjd", scale="tai"), mode
1237
        )
1238

1239
        # determining which scheme to use to filter slews
1240
        obsModName = Obs.__class__.__name__
×
1241

1242
        # slew times have not been calculated/decided yet (SotoStarshade)
1243
        if obsModName == "SotoStarshade":
×
1244
            sInds, intTimes, slewTimes, dV = self.findAllowableOcculterSlews(
×
1245
                sInds,
1246
                old_sInd,
1247
                sd[sInds],
1248
                slewTimes[sInds],
1249
                obsTimeArray[sInds, :],
1250
                intTimeArray[sInds, :],
1251
                mode,
1252
            )
1253

1254
        # slew times were calculated/decided beforehand (Observatory Prototype)
1255
        else:
1256
            sInds, intTimes, slewTimes = self.filterOcculterSlews(
×
1257
                sInds,
1258
                slewTimes[sInds],
1259
                obsTimeArray[sInds, :],
1260
                intTimeArray[sInds, :],
1261
                mode,
1262
            )
1263
            dV = np.zeros(len(sInds)) * u.m / u.s
×
1264

1265
        return sInds, slewTimes, intTimes, dV
×
1266

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

1270
        Used by the refineOcculterSlews method when slew times have been selected
1271
        a priori. This method filters out slews that are not within desired observing
1272
        blocks, the maximum allowed integration time, and are outside of future
1273
        keepouts.
1274

1275
        Args:
1276
            sInds (~numpy.ndarray(int)):
1277
                Indices of available targets
1278
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1279
                slew times to all stars (must be indexed by sInds)
1280
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1281
                Array of times during which a star is out of keepout, has shape
1282
                nx50 where n is the number of stars in sInds. Unit of days
1283
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1284
                Array of integration times for each time in obsTimeArray, has shape
1285
                nx2 where n is the number of stars in sInds. Unit of days
1286
            mode (dict):
1287
                Selected observing mode for detection
1288

1289
        Returns:
1290
            tuple:
1291
                sInds (int):
1292
                    Indeces of next target star
1293
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1294
                    Integration times for detection in units of day
1295
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1296
                    slew times to all stars (must be indexed by sInds)
1297
        """
1298

1299
        TK = self.TimeKeeping
×
1300
        Obs = self.Observatory
×
1301

1302
        # allocate settling time + overhead time
1303
        tmpCurrentTimeAbs = (
×
1304
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1305
        )
1306
        tmpCurrentTimeNorm = (
×
1307
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1308
        )
1309

1310
        # 0. lambda function that linearly interpolates Integration Time
1311
        # between obsTimes
1312
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1313
            t - np.array(x[:, 0]).reshape(len(t), 1)
1314
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1315

1316
        # 1. initializing arrays
1317
        obsTimesRange = np.array(
×
1318
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1319
        )  # nx2 array with start and end times of obsTimes for each star
1320
        intTimesRange = np.array([intTimeArray[:, 0], intTimeArray[:, -1]])
×
1321

1322
        OBnumbers = np.zeros(
×
1323
            [len(sInds), 1]
1324
        )  # for each sInd, will show during which OB observations will take place
1325
        maxIntTimes = np.zeros([len(sInds), 1]) * u.d
×
1326

1327
        intTimes = (
×
1328
            linearInterp(
1329
                intTimesRange.T,
1330
                obsTimesRange.T,
1331
                (tmpCurrentTimeAbs + slewTimes).reshape(len(sInds), 1).value,
1332
            )
1333
            * u.d
1334
        )  # calculate intTimes for each slew time
1335

1336
        minObsTimeNorm = (obsTimesRange[0, :] - tmpCurrentTimeAbs.value).reshape(
×
1337
            [len(sInds), 1]
1338
        )
1339
        maxObsTimeNorm = (obsTimesRange[1, :] - tmpCurrentTimeAbs.value).reshape(
×
1340
            [len(sInds), 1]
1341
        )
1342
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1343

1344
        # 2. find OBnumber for each sInd's slew time
1345
        if len(TK.OBendTimes) > 1:
×
1346
            for i in range(len(sInds)):
×
1347
                S = np.where(
×
1348
                    TK.OBstartTimes.value - tmpCurrentTimeNorm.value
1349
                    < slewTimes[i].value
1350
                )[0][-1]
1351
                F = np.where(
×
1352
                    TK.OBendTimes.value - tmpCurrentTimeNorm.value < slewTimes[i].value
1353
                )[0]
1354

1355
                # case when slews are in the first OB
1356
                if F.shape[0] == 0:
×
1357
                    F = -1
×
1358
                else:
1359
                    F = F[-1]
×
1360

1361
                # slew occurs within an OB (nth OB has started but hasn't ended)
1362
                if S != F:
×
1363
                    OBnumbers[i] = S
×
1364
                    (
×
1365
                        maxIntTimeOBendTime,
1366
                        maxIntTimeExoplanetObsTime,
1367
                        maxIntTimeMissionLife,
1368
                    ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, TK.OBstartTimes[S], S)
1369
                    maxIntTimes[i] = min(
×
1370
                        maxIntTimeOBendTime,
1371
                        maxIntTimeExoplanetObsTime,
1372
                        maxIntTimeMissionLife,
1373
                    )  # Maximum intTime allowed
1374

1375
                # slew occurs between OBs, badbadnotgood
1376
                else:
1377
                    OBnumbers[i] = -1
×
1378
                    maxIntTimes[i] = 0 * u.d
×
1379
            OBstartTimeNorm = (
×
1380
                TK.OBstartTimes[np.array(OBnumbers, dtype=int)].value
1381
                - tmpCurrentTimeNorm.value
1382
            )
1383
        else:
1384
            (
×
1385
                maxIntTimeOBendTime,
1386
                maxIntTimeExoplanetObsTime,
1387
                maxIntTimeMissionLife,
1388
            ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm)
1389
            maxIntTimes[:] = min(
×
1390
                maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1391
            )  # Maximum intTime allowed
1392
            OBstartTimeNorm = np.zeros(OBnumbers.shape)
×
1393

1394
        # finding the minimum possible slew time
1395
        # (either OBstart or when star JUST leaves keepout)
1396
        minAllowedSlewTime = np.max([OBstartTimeNorm, minObsTimeNorm], axis=0)
×
1397

1398
        # 3. start filtering process
1399
        good_inds = np.where((OBnumbers >= 0) & (ObsTimeRange > intTimes.value))[0]
×
1400
        # star slew within OB -AND- can finish observing
1401
        # before it goes back into keepout
1402
        if good_inds.shape[0] > 0:
×
1403
            # the good ones
1404
            sInds = sInds[good_inds]
×
1405
            slewTimes = slewTimes[good_inds]
×
1406
            intTimes = intTimes[good_inds]
×
1407
            OBstartTimeNorm = OBstartTimeNorm[good_inds]
×
1408
            minAllowedSlewTime = minAllowedSlewTime[good_inds]
×
1409

1410
            # maximum allowed slew time based on integration times
1411
            maxAllowedSlewTime = maxIntTimes[good_inds].value - intTimes.value
×
1412
            maxAllowedSlewTime[maxAllowedSlewTime < 0] = -np.inf
×
1413
            maxAllowedSlewTime += OBstartTimeNorm  # calculated rel to currentTime norm
×
1414

1415
            # checking to see if slewTimes are allowed
1416
            good_inds = np.where(
×
1417
                (slewTimes.reshape([len(sInds), 1]).value > minAllowedSlewTime)
1418
                & (slewTimes.reshape([len(sInds), 1]).value < maxAllowedSlewTime)
1419
            )[0]
1420

1421
            slewTimes = slewTimes[good_inds]
×
1422
        else:
1423
            slewTimes = slewTimes[good_inds]
×
1424

1425
        return sInds[good_inds], intTimes[good_inds].flatten(), slewTimes
×
1426

1427
    def findAllowableOcculterSlews(
1✔
1428
        self, sInds, old_sInd, sd, slewTimes, obsTimeArray, intTimeArray, mode
1429
    ):
1430
        """Finds an array of allowable slew times for each star
1431

1432
        Used by the refineOcculterSlews method when slew times have NOT been selected
1433
        a priori. This method creates nx50 arrays (where the row corresponds to a
1434
        specific star and the column corresponds to a future point in time relative to
1435
        currentTime).
1436

1437
        These arrays are initially zero but are populated with the corresponding values
1438
        (slews, intTimes, etc) if slewing to that time point (i.e. beginning an
1439
        observation) would lead to a successful observation. A "successful observation"
1440
        is defined by certain conditions relating to keepout and the komap, observing
1441
        blocks, mission lifetime, and some constraints on the dVmap calculation in
1442
        SotoStarshade. Each star will likely have a range of slewTimes that would lead
1443
        to a successful observation -- another method is then called to select the best
1444
        of these slewTimes.
1445

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

1464
        Returns:
1465
            tuple:
1466
                sInds (numpy.ndarray(int)):
1467
                    Indices of next target star
1468
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1469
                    slew times to all stars (must be indexed by sInds)
1470
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1471
                    Integration times for detection in units of day
1472
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1473
                    Delta-V used to transfer to new star line of sight in unis of m/s
1474
        """
1475
        TK = self.TimeKeeping
×
1476
        Obs = self.Observatory
×
1477
        TL = self.TargetList
×
1478

1479
        # 0. lambda function that linearly interpolates Integration
1480
        # Time between obsTimes
1481
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1482
            t - np.array(x[:, 0]).reshape(len(t), 1)
1483
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1484

1485
        # allocate settling time + overhead time
1486
        tmpCurrentTimeAbs = (
×
1487
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1488
        )
1489
        tmpCurrentTimeNorm = (
×
1490
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1491
        )
1492

1493
        # 1. initializing arrays
1494
        obsTimes = np.array(
×
1495
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1496
        )  # nx2 array with start and end times of obsTimes for each star
1497
        intTimes_int = (
×
1498
            np.zeros(obsTimeArray.shape) * u.d
1499
        )  # initializing intTimes of shape nx50 then interpolating
1500
        intTimes_int = (
×
1501
            np.hstack(
1502
                [
1503
                    intTimeArray[:, 0].reshape(len(sInds), 1).value,
1504
                    linearInterp(
1505
                        intTimeArray.value, obsTimes.T, obsTimeArray[:, 1:].value
1506
                    ),
1507
                ]
1508
            )
1509
            * u.d
1510
        )
1511
        allowedSlewTimes = np.zeros(obsTimeArray.shape) * u.d
×
1512
        allowedintTimes = np.zeros(obsTimeArray.shape) * u.d
×
1513
        allowedCharTimes = np.zeros(obsTimeArray.shape) * u.d
×
1514
        obsTimeArrayNorm = obsTimeArray.value - tmpCurrentTimeAbs.value
×
1515

1516
        # obsTimes -> relative to current Time
1517
        try:
×
1518
            minObsTimeNorm = np.array([np.min(v[v > 0]) for v in obsTimeArrayNorm])
×
1519
        except:  # noqa: E722
×
1520
            # an error pops up sometimes at the end of the mission, this fixes it
1521
            # TODO: define the error type that occurs
1522
            # rewrite to avoid a try/except if possible
1523
            minObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1524

1525
        maxObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1526
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1527

1528
        # getting max possible intTime
1529
        (
×
1530
            maxIntTimeOBendTime,
1531
            maxIntTimeExoplanetObsTime,
1532
            maxIntTimeMissionLife,
1533
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm, TK.OBnumber)
1534
        maxIntTime = min(
×
1535
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1536
        )  # Maximum intTime allowed
1537

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

1542
        # first filled in for the current OB
1543
        minAllowedSlewTimes = np.array(
×
1544
            [minObsTimeNorm.T] * len(intTimes_int.T)
1545
        ).T  # just to make it nx50 so it plays nice with the other arrays
1546
        maxAllowedSlewTimes = maxIntTime.value - intTimes_int.value
×
1547
        maxAllowedSlewTimes[maxAllowedSlewTimes > Obs.occ_dtmax.value] = (
×
1548
            Obs.occ_dtmax.value
1549
        )
1550

1551
        # conditions that must be met to define an allowable slew time
1552
        cond1 = (
×
1553
            minAllowedSlewTimes >= Obs.occ_dtmin.value
1554
        )  # minimum dt time in dV map interpolant
1555
        cond2 = (
×
1556
            maxAllowedSlewTimes <= Obs.occ_dtmax.value
1557
        )  # maximum dt time in dV map interpolant
1558
        cond3 = maxAllowedSlewTimes > minAllowedSlewTimes
×
1559
        cond4 = intTimes_int.value < ObsTimeRange.reshape(len(sInds), 1)
×
1560

1561
        conds = cond1 & cond2 & cond3 & cond4
×
1562
        minAllowedSlewTimes[np.invert(conds)] = (
×
1563
            np.inf
1564
        )  # these are filtered during the next filter
1565
        maxAllowedSlewTimes[np.invert(conds)] = -np.inf
×
1566

1567
        # one last condition to meet
1568
        map_i, map_j = np.where(
×
1569
            (obsTimeArrayNorm > minAllowedSlewTimes)
1570
            & (obsTimeArrayNorm < maxAllowedSlewTimes)
1571
        )
1572

1573
        # 2.5 if any stars are slew-able to within this OB block, populate
1574
        # "allowedSlewTimes", a running tally of possible slews
1575
        # within the time range a star is observable (out of keepout)
1576
        if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1577
            allowedSlewTimes[map_i, map_j] = obsTimeArrayNorm[map_i, map_j] * u.d
×
1578
            allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1579
            allowedCharTimes[map_i, map_j] = maxIntTime - intTimes_int[map_i, map_j]
×
1580

1581
        # 3. search future OBs
1582
        OB_withObsStars = (
×
1583
            TK.OBstartTimes.value - np.min(obsTimeArrayNorm) - tmpCurrentTimeNorm.value
1584
        )  # OBs within which any star is observable
1585

1586
        if any(OB_withObsStars > 0):
×
1587
            nOBstart = np.argmin(np.abs(OB_withObsStars))
×
1588
            nOBend = np.argmax(OB_withObsStars)
×
1589

1590
            # loop through the next 5 OBs (or until mission is over if there are less
1591
            # than 5 OBs in the future)
1592
            for i in np.arange(nOBstart, np.min([nOBend, nOBstart + 5])):
×
1593
                # max int Times for the next OB
UNCOV
1594
                (
×
1595
                    maxIntTimeOBendTime,
1596
                    maxIntTimeExoplanetObsTime,
1597
                    maxIntTimeMissionLife,
1598
                ) = TK.get_ObsDetectionMaxIntTime(
1599
                    Obs, mode, TK.OBstartTimes[i + 1], i + 1
1600
                )
1601
                maxIntTime_nOB = min(
×
1602
                    maxIntTimeOBendTime,
1603
                    maxIntTimeExoplanetObsTime,
1604
                    maxIntTimeMissionLife,
1605
                )  # Maximum intTime allowed
1606

1607
                # min and max slew times rel to current Time (norm)
1608
                nOBstartTimeNorm = np.array(
×
1609
                    [TK.OBstartTimes[i + 1].value - tmpCurrentTimeNorm.value]
1610
                    * len(sInds)
1611
                )
1612

1613
                # min slew times for stars start either whenever the star first leaves
1614
                # keepout or when next OB stars, whichever comes last
1615
                minAllowedSlewTimes_nOB = np.array(
×
1616
                    [np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T]
1617
                    * len(maxAllowedSlewTimes.T)
1618
                ).T
1619
                maxAllowedSlewTimes_nOB = (
×
1620
                    nOBstartTimeNorm.reshape(len(sInds), 1)
1621
                    + maxIntTime_nOB.value
1622
                    - intTimes_int.value
1623
                )
1624
                maxAllowedSlewTimes_nOB[
×
1625
                    maxAllowedSlewTimes_nOB > Obs.occ_dtmax.value
1626
                ] = Obs.occ_dtmax.value
1627

1628
                # amount of time left when the stars are still out of keepout
1629
                ObsTimeRange_nOB = (
×
1630
                    maxObsTimeNorm
1631
                    - np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T
1632
                )
1633

1634
                # condition to be met for an allowable slew time
1635
                cond1 = minAllowedSlewTimes_nOB >= Obs.occ_dtmin.value
×
1636
                cond2 = maxAllowedSlewTimes_nOB <= Obs.occ_dtmax.value
×
1637
                cond3 = maxAllowedSlewTimes_nOB > minAllowedSlewTimes_nOB
×
1638
                cond4 = intTimes_int.value < ObsTimeRange_nOB.reshape(len(sInds), 1)
×
1639
                cond5 = intTimes_int.value < maxIntTime_nOB.value
×
1640
                conds = cond1 & cond2 & cond3 & cond4 & cond5
×
1641

1642
                minAllowedSlewTimes_nOB[np.invert(conds)] = np.inf
×
1643
                maxAllowedSlewTimes_nOB[np.invert(conds)] = -np.inf
×
1644

1645
                # one last condition
1646
                map_i, map_j = np.where(
×
1647
                    (obsTimeArrayNorm > minAllowedSlewTimes_nOB)
1648
                    & (obsTimeArrayNorm < maxAllowedSlewTimes_nOB)
1649
                )
1650

1651
                # 3.33 populate the running tally of allowable slew times if it meets
1652
                # all conditions
1653
                if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1654
                    allowedSlewTimes[map_i, map_j] = (
×
1655
                        obsTimeArrayNorm[map_i, map_j] * u.d
1656
                    )
1657
                    allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1658
                    allowedCharTimes[map_i, map_j] = (
×
1659
                        maxIntTime_nOB - intTimes_int[map_i, map_j]
1660
                    )
1661

1662
        # 3.67 filter out any stars that are not observable at all
1663
        filterDuds = np.sum(allowedSlewTimes, axis=1) > 0.0
×
1664
        sInds = sInds[filterDuds]
×
1665

1666
        # 4. choose a slew time for each available star
1667
        # calculate dVs for each possible slew time for each star
1668
        allowed_dVs = Obs.calculate_dV(
×
1669
            TL,
1670
            old_sInd,
1671
            sInds,
1672
            sd[filterDuds],
1673
            allowedSlewTimes[filterDuds, :],
1674
            tmpCurrentTimeAbs,
1675
        )
1676

1677
        if len(sInds.tolist()) > 0:
×
1678
            # select slew time for each star
1679
            dV_inds = np.arange(0, len(sInds))
×
1680
            sInds, intTime, slewTime, dV = self.chooseOcculterSlewTimes(
×
1681
                sInds,
1682
                allowedSlewTimes[filterDuds, :],
1683
                allowed_dVs[dV_inds, :],
1684
                allowedintTimes[filterDuds, :],
1685
                allowedCharTimes[filterDuds, :],
1686
            )
1687

1688
            return sInds, intTime, slewTime, dV
×
1689

1690
        else:
1691
            empty = np.asarray([], dtype=int)
×
1692
            return empty, empty * u.d, empty * u.d, empty * u.m / u.s
×
1693

1694
    def chooseOcculterSlewTimes(self, sInds, slewTimes, dV, intTimes, charTimes):
1✔
1695
        """Selects the best slew time for each star
1696

1697
        This method searches through an array of permissible slew times for
1698
        each star and chooses the best slew time for the occulter based on
1699
        maximizing possible characterization time for that particular star (as
1700
        a default).
1701

1702
        Args:
1703
            sInds (~numpy.ndarray(int)):
1704
                Indices of available targets
1705
            slewTimes (astropy quantity array):
1706
                slew times to all stars (must be indexed by sInds)
1707
            dV (~astropy.units.Quantity(~numpy.ndarray(float))):
1708
                Delta-V used to transfer to new star line of sight in unis of m/s
1709
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1710
                Integration times for detection in units of day
1711
            charTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1712
                Time left over after integration which could be used for
1713
                characterization in units of day
1714

1715
        Returns:
1716
            tuple:
1717
                sInds (int):
1718
                    Indeces of next target star
1719
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1720
                    slew times to all stars (must be indexed by sInds)
1721
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1722
                    Integration times for detection in units of day
1723
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1724
                    Delta-V used to transfer to new star line of sight in unis of m/s
1725
        """
1726

1727
        # selection criteria for each star slew
1728
        good_j = np.argmax(
×
1729
            charTimes, axis=1
1730
        )  # maximum possible characterization time available
1731
        good_i = np.arange(0, len(sInds))
×
1732

1733
        dV = dV[good_i, good_j]
×
1734
        intTime = intTimes[good_i, good_j]
×
1735
        slewTime = slewTimes[good_i, good_j]
×
1736

1737
        return sInds, intTime, slewTime, dV
×
1738

1739
    def observation_detection(self, sInd, intTime, mode):
1✔
1740
        """Determines SNR and detection status for a given integration time
1741
        for detection. Also updates the lastDetected and starRevisit lists.
1742

1743
        Args:
1744
            sInd (int):
1745
                Integer index of the star of interest
1746
            intTime (~astropy.units.Quantity(~numpy.ndarray(float))):
1747
                Selected star integration time for detection in units of day.
1748
                Defaults to None.
1749
            mode (dict):
1750
                Selected observing mode for detection
1751

1752
        Returns:
1753
            tuple:
1754
                detected (numpy.ndarray(int)):
1755
                    Detection status for each planet orbiting the observed target star:
1756
                    1 is detection, 0 missed detection, -1 below IWA, and -2 beyond OWA
1757
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
1758
                    Surface brightness of local zodiacal light in units of 1/arcsec2
1759
                JEZ (astropy.units.Quantity(numpy.ndarray(float))):
1760
                    Intensity of exo-zodiacal light in units of photons/s/m2/arcsec2
1761
                systemParams (dict):
1762
                    Dictionary of time-dependant planet properties averaged over the
1763
                    duration of the integration
1764
                SNR (numpy.darray(float)):
1765
                    Detection signal-to-noise ratio of the observable planets
1766
                FA (bool):
1767
                    False alarm (false positive) boolean
1768

1769
        """
1770

1771
        PPop = self.PlanetPopulation
1✔
1772
        ZL = self.ZodiacalLight
1✔
1773
        PPro = self.PostProcessing
1✔
1774
        TL = self.TargetList
1✔
1775
        SU = self.SimulatedUniverse
1✔
1776
        Obs = self.Observatory
1✔
1777
        TK = self.TimeKeeping
1✔
1778

1779
        # Save Current Time before attempting time allocation
1780
        currentTimeNorm = TK.currentTimeNorm.copy()
1✔
1781
        currentTimeAbs = TK.currentTimeAbs.copy()
1✔
1782

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

1796
        # initialize outputs
1797
        detected = np.array([], dtype=int)
1✔
1798
        # write current system params by default
1799
        systemParams = SU.dump_system_params(sInd)
1✔
1800
        SNR = np.zeros(len(pInds))
1✔
1801

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

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

1849
        # if no planet, just save zodiacal brightness in the middle of the integration
1850
        else:
1851
            totTime = intTime * (mode["timeMultiplier"])
1✔
1852
            fZ = ZL.fZ(Obs, TL, sInd, currentTimeAbs + totTime / 2.0, mode)[0]
1✔
1853
            # Use the default star value if no planets
1854
            JEZ = TL.JEZ0[mode["hex"]][sInd]
1✔
1855

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

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

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

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

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

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

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

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

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

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

1940
        Updates self.starRevisit attribute only
1941

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

1952
        Returns:
1953
            None
1954

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

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

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

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

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

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

2025
        """
2026

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2173
            pIndsChar = pIndsDet[tochar]
1✔
2174
            log_char = "   - Charact. planet inds %s (%s/%s detected)" % (
1✔
2175
                pIndsChar,
2176
                len(pIndsChar),
2177
                len(pIndsDet),
2178
            )
2179
            self.logger.info(log_char)
1✔
2180
            self.vprint(log_char)
1✔
2181

2182
            # SNR CALCULATION:
2183
            # first, calculate SNR for observable planets (without false alarm)
2184
            planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
1✔
2185
            SNRplans = np.zeros(len(planinds))
1✔
2186
            if len(planinds) > 0:
1✔
2187
                # initialize arrays for SNR integration
2188
                fZs = np.zeros(self.ntFlux) / u.arcsec**2.0
1✔
2189
                systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
2190
                JEZs = (
1✔
2191
                    np.zeros((self.ntFlux, len(planinds)))
2192
                    * u.ph
2193
                    / u.s
2194
                    / u.m**2
2195
                    / u.arcsec**2
2196
                )
2197
                Ss = np.zeros((self.ntFlux, len(planinds)))
1✔
2198
                Ns = np.zeros((self.ntFlux, len(planinds)))
1✔
2199
                # integrate the signal (planet flux) and noise
2200
                dt = intTime / float(self.ntFlux)
1✔
2201
                timePlus = (
1✔
2202
                    Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2203
                )  # accounts for the time since the current time
2204
                for i in range(self.ntFlux):
1✔
2205
                    # allocate first half of dt
2206
                    timePlus += dt / 2.0
1✔
2207
                    # calculate current zodiacal light brightness
2208
                    fZs[i] = ZL.fZ(Obs, TL, sInd, currentTimeAbs + timePlus, mode)[0]
1✔
2209
                    # propagate the system to match up with current time
2210
                    SU.propag_system(
1✔
2211
                        sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2212
                    )
2213
                    self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
2214
                    # Calculate the exozodi intensity
2215
                    JEZs[i] = SU.scale_JEZ(sInd, mode, pInds=planinds)
1✔
2216
                    # save planet parameters
2217
                    systemParamss[i] = SU.dump_system_params(sInd)
1✔
2218
                    # calculate signal and noise (electron count rates)
2219
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
2220
                        sInd, planinds, dt, mode, fZ=fZs[i], JEZ=JEZs[i]
2221
                    )
2222
                    # allocate second half of dt
2223
                    timePlus += dt / 2.0
1✔
2224

2225
                # average output parameters
2226
                fZ = np.mean(fZs)
1✔
2227
                JEZ = np.mean(JEZs, axis=0)
1✔
2228
                systemParams = {
1✔
2229
                    key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2230
                    / float(self.ntFlux)
2231
                    for key in sorted(systemParamss[0])
2232
                }
2233
                # calculate planets SNR
2234
                S = Ss.sum(0)
1✔
2235
                N = Ns.sum(0)
1✔
2236
                SNRplans[N > 0] = S[N > 0] / N[N > 0]
1✔
2237
                # allocate extra time for timeMultiplier
2238

2239
            # if only a FA, just save zodiacal brightness in the middle of the
2240
            # integration
2241
            else:
2242
                totTime = intTime * (mode["timeMultiplier"])
×
2243
                fZ = ZL.fZ(Obs, TL, sInd, currentTimeAbs + totTime / 2.0, mode)[0]
×
2244
                # Use the default star value if no planets
NEW
2245
                JEZ = TL.JEZ0[mode["hex"]][sInd]
×
2246

2247
            # calculate the false alarm SNR (if any)
2248
            SNRfa = []
1✔
2249
            if pIndsChar[-1] == -1:
1✔
NEW
2250
                JEZ = self.lastDetected[sInd, 1][-1]
×
2251
                dMag = self.lastDetected[sInd, 2][-1]
×
2252
                WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
NEW
2253
                C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, JEZ, dMag, WA, mode, TK=TK)
×
2254
                S = (C_p * intTime).decompose().value
×
2255
                N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2256
                SNRfa = S / N if N > 0.0 else 0.0
×
2257

2258
            # save all SNRs (planets and FA) to one array
2259
            SNRinds = np.where(det)[0][tochar]
1✔
2260
            SNR[SNRinds] = np.append(SNRplans, SNRfa)
1✔
2261

2262
            # now, store characterization status: 1 for full spectrum,
2263
            # -1 for partial spectrum, 0 for not characterized
2264
            char = SNR >= mode["SNR"]
1✔
2265
            # initialize with full spectra
2266
            characterized = char.astype(int)
1✔
2267
            WAchar = self.lastDetected[sInd, 3][char] * u.arcsec
1✔
2268
            # find the current WAs of characterized planets
2269
            WAs = systemParams["WA"]
1✔
2270
            if FA:
1✔
2271
                WAs = np.append(WAs, self.lastDetected[sInd, 3][-1] * u.arcsec)
×
2272
            # check for partial spectra (for coronagraphs only)
2273
            if not (mode["syst"]["occulter"]):
1✔
2274
                IWA_max = mode["IWA"] * (1.0 + mode["BW"] / 2.0)
1✔
2275
                OWA_min = mode["OWA"] * (1.0 - mode["BW"] / 2.0)
1✔
2276
                char[char] = (WAchar < IWA_max) | (WAchar > OWA_min)
1✔
2277
                characterized[char] = -1
1✔
2278
            # encode results in spectra lists (only for planets, not FA)
2279
            charplans = characterized[:-1] if FA else characterized
1✔
2280
            self.fullSpectra[pInds[charplans == 1]] += 1
1✔
2281
            self.partialSpectra[pInds[charplans == -1]] += 1
1✔
2282

2283
        if self.make_debug_bird_plots:
1✔
2284
            from tools.obs_plot import obs_plot
×
2285

2286
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2287

2288
        return characterized.astype(int), fZ, JEZ, systemParams, SNR, intTime
1✔
2289

2290
    def calc_signal_noise(
1✔
2291
        self, sInd, pInds, t_int, mode, fZ=None, JEZ=None, dMag=None, WA=None
2292
    ):
2293
        """Calculates the signal and noise fluxes for a given time interval. Called
2294
        by observation_detection and observation_characterization methods in the
2295
        SurveySimulation module.
2296

2297
        Args:
2298
            sInd (int):
2299
                Integer index of the star of interest
2300
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2301
                Integration time interval in units of day
2302
            pInds (int):
2303
                Integer indices of the planets of interest
2304
            mode (dict):
2305
                Selected observing mode (from OpticalSystem)
2306
            fZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2307
                Surface brightness of local zodiacal light in units of 1/arcsec2
2308
            JEZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2309
                Intensity of exo-zodiacal light in units of ph/s/m2/arcsec2
2310
            dMag (~numpy.ndarray(float)):
2311
                Differences in magnitude between planets and their host star
2312
            WA (~astropy.units.Quantity(~numpy.ndarray(float))):
2313
                Working angles of the planets of interest in units of arcsec
2314

2315
        Returns:
2316
            tuple:
2317
                Signal (float):
2318
                    Counts of signal
2319
                Noise (float):
2320
                    Counts of background noise variance
2321

2322
        """
2323

2324
        OS = self.OpticalSystem
1✔
2325
        ZL = self.ZodiacalLight
1✔
2326
        TL = self.TargetList
1✔
2327
        SU = self.SimulatedUniverse
1✔
2328
        Obs = self.Observatory
1✔
2329
        TK = self.TimeKeeping
1✔
2330

2331
        # calculate optional parameters if not provided
2332
        fZ = (
1✔
2333
            fZ
2334
            if (fZ is not None)
2335
            else ZL.fZ(Obs, TL, sInd, TK.currentTimeAbs.copy(), mode)
2336
        )
2337
        JEZ = (
1✔
2338
            u.Quantity(JEZ, ndmin=1)
2339
            if (JEZ is not None)
2340
            else SU.scale_JEZ(sInd, mode, pInds=pInds)
2341
        )
2342

2343
        # if lucky_planets, use lucky planet params for dMag and WA
2344
        if SU.lucky_planets and mode in list(
1✔
2345
            filter(lambda mode: "spec" in mode["inst"]["name"], OS.observingModes)
2346
        ):
2347
            phi = (1 / np.pi) * np.ones(len(SU.d))
×
2348
            dMag = deltaMag(SU.p, SU.Rp, SU.d, phi)[pInds]  # delta magnitude
×
2349
            WA = np.arctan(SU.a / TL.dist[SU.plan2star]).to("arcsec")[
×
2350
                pInds
2351
            ]  # working angle
2352
        else:
2353
            dMag = dMag if (dMag is not None) else SU.dMag[pInds]
1✔
2354
            WA = WA if (WA is not None) else SU.WA[pInds]
1✔
2355

2356
        # initialize Signal and Noise arrays
2357
        Signal = np.zeros(len(pInds))
1✔
2358
        Noise = np.zeros(len(pInds))
1✔
2359

2360
        # find observable planets wrt IWA-OWA range
2361
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2362

2363
        if np.any(obs):
1✔
2364
            # find electron counts
2365
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(
1✔
2366
                TL, sInd, fZ, JEZ[obs], dMag[obs], WA[obs], mode, TK=TK
2367
            )
2368
            # calculate signal and noise levels (based on Nemati14 formula)
2369
            Signal[obs] = (C_p * t_int).decompose().value
1✔
2370
            Noise[obs] = np.sqrt((C_b * t_int + (C_sp * t_int) ** 2).decompose().value)
1✔
2371

2372
        return Signal, Noise
1✔
2373

2374
    def update_occulter_mass(self, DRM, sInd, t_int, skMode):
1✔
2375
        """Updates the occulter wet mass in the Observatory module, and stores all
2376
        the occulter related values in the DRM array.
2377

2378
        Args:
2379
            DRM (dict):
2380
                Design Reference Mission, contains the results of one complete
2381
                observation (detection and characterization)
2382
            sInd (int):
2383
                Integer index of the star of interest
2384
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2385
                Selected star integration time (for detection or characterization)
2386
                in units of day
2387
            skMode (str):
2388
                Station keeping observing mode type ('det' or 'char')
2389

2390
        Returns:
2391
            dict:
2392
                Design Reference Mission dictionary, contains the results of one
2393
                complete observation
2394

2395
        """
2396

2397
        TL = self.TargetList
×
2398
        Obs = self.Observatory
×
2399
        TK = self.TimeKeeping
×
2400

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

2403
        # decrement mass for station-keeping
2404
        dF_lateral, dF_axial, intMdot, mass_used, deltaV = Obs.mass_dec_sk(
×
2405
            TL, sInd, TK.currentTimeAbs.copy(), t_int
2406
        )
2407

2408
        DRM[skMode + "_dV"] = deltaV.to("m/s")
×
2409
        DRM[skMode + "_mass_used"] = mass_used.to("kg")
×
2410
        DRM[skMode + "_dF_lateral"] = dF_lateral.to("N")
×
2411
        DRM[skMode + "_dF_axial"] = dF_axial.to("N")
×
2412
        # update spacecraft mass
2413
        Obs.scMass = Obs.scMass - mass_used
×
2414
        DRM["scMass"] = Obs.scMass.to("kg")
×
2415
        if Obs.twotanks:
×
2416
            Obs.skMass = Obs.skMass - mass_used
×
2417
            DRM["skMass"] = Obs.skMass.to("kg")
×
2418

2419
        return DRM
×
2420

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

2424
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2425
        objects with their own outspecs.
2426

2427
        Args:
2428
            genNewPlanets (bool):
2429
                Generate all new planets based on the original input specification.
2430
                If False, then the original planets will remain. Setting to True forces
2431
                ``rewindPlanets`` to be True as well. Defaults True.
2432
            rewindPlanets (bool):
2433
                Reset the current set of planets to their original orbital phases.
2434
                If both genNewPlanets and rewindPlanet  are False, the original planets
2435
                will be retained and will not be rewound to their initial starting
2436
                locations (i.e., all systems will remain at the times they were at the
2437
                end of the last run, thereby effectively randomizing planet phases.
2438
                Defaults True.
2439
            seed (int, optional):
2440
                Random seed to use for all random number generation. If None (default)
2441
                a new random seed will be generated when re-initializing the
2442
                SurveySimulation.
2443

2444
        """
2445

2446
        SU = self.SimulatedUniverse
1✔
2447
        TK = self.TimeKeeping
1✔
2448
        TL = self.TargetList
1✔
2449

2450
        # re-initialize SurveySimulation arrays
2451
        specs = self._outspec
1✔
2452
        specs["modules"] = self.modules
1✔
2453

2454
        if seed is None:  # pop the seed so a new one is generated
1✔
2455
            if "seed" in specs:
1✔
2456
                specs.pop("seed")
1✔
2457
        else:  # if seed is provided, replace seed with provided seed
2458
            specs["seed"] = seed
×
2459

2460
        # reset mission time, re-init surveysim and observatory
2461
        TK.__init__(**TK._outspec)
1✔
2462
        self.__init__(**specs)
1✔
2463
        self.Observatory.__init__(**self.Observatory._outspec)
1✔
2464

2465
        # generate new planets if requested (default)
2466
        if genNewPlanets:
1✔
2467
            TL.stellar_mass()
1✔
2468
            TL.I = TL.gen_inclinations(TL.PlanetPopulation.Irange)  # noqa: E741
1✔
2469
            SU.gen_physical_properties(**SU._outspec)
1✔
2470
            rewindPlanets = True
1✔
2471
        # re-initialize systems if requested (default)
2472
        if rewindPlanets:
1✔
2473
            SU.init_systems()
1✔
2474

2475
        # reset helper arrays
2476
        self.initializeStorageArrays()
1✔
2477

2478
        self.vprint("Simulation reset.")
1✔
2479

2480
    def genOutSpec(
1✔
2481
        self,
2482
        tofile: Optional[str] = None,
2483
        starting_outspec: Optional[Dict[str, Any]] = None,
2484
        modnames: bool = False,
2485
    ) -> Dict[str, Any]:
2486
        """Join all _outspec dicts from all modules into one output dict
2487
        and optionally write out to JSON file on disk.
2488

2489
        Args:
2490
            tofile (str, optional):
2491
                Name of the file containing all output specifications (outspecs).
2492
                Defaults to None.
2493
            starting_outspec (dict, optional):
2494
                Initial outspec (from MissionSim). Defaults to None.
2495
            modnames (bool):
2496
                If True, populate outspec dictionary with the module it originated from,
2497
                instead of the actual value of the keyword. Defaults False.
2498

2499
        Returns:
2500
            dict:
2501
                Dictionary containing the full :ref:`sec:inputspec`, including all
2502
                filled-in default values. Combination of all individual module _outspec
2503
                attributes.
2504

2505
        """
2506

2507
        # start with our own outspec
2508
        if modnames:
1✔
2509
            out = copy.copy(self._outspec)
1✔
2510
            for k in out:
1✔
2511
                out[k] = self.__class__.__name__
1✔
2512
        else:
2513
            out = copy.deepcopy(self._outspec)
1✔
2514

2515
        # Add any provided other outspec
2516
        if starting_outspec is not None:
1✔
2517
            out.update(starting_outspec)
1✔
2518

2519
        # add in all modules _outspec's
2520
        for module in self.modules.values():
1✔
2521
            if modnames:
1✔
2522
                tmp = copy.copy(module._outspec)
1✔
2523
                for k in tmp:
1✔
2524
                    tmp[k] = module.__class__.__name__
1✔
2525
            else:
2526
                tmp = module._outspec
1✔
2527
            out.update(tmp)
1✔
2528

2529
        # add in the specific module names used
2530
        out["modules"] = {}
1✔
2531
        for mod_name, module in self.modules.items():
1✔
2532
            # find the module file
2533
            mod_name_full = module.__module__
1✔
2534
            if mod_name_full.startswith("EXOSIMS"):
1✔
2535
                # take just its short name if it is in EXOSIMS
2536
                mod_name_short = mod_name_full.split(".")[-1]
1✔
2537
            else:
2538
                # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2539
                mod_name_short = re.sub(
×
2540
                    r"\.pyc$", ".py", inspect.getfile(module.__class__)
2541
                )
2542
            out["modules"][mod_name] = mod_name_short
1✔
2543
        # add catalog name
2544
        if self.TargetList.keepStarCatalog:
1✔
2545
            module = self.TargetList.StarCatalog
×
2546
            mod_name_full = module.__module__
×
2547
            if mod_name_full.startswith("EXOSIMS"):
×
2548
                # take just its short name if it is in EXOSIMS
2549
                mod_name_short = mod_name_full.split(".")[-1]
×
2550
            else:
2551
                # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2552
                mod_name_short = re.sub(
×
2553
                    r"\.pyc$", ".py", inspect.getfile(module.__class__)
2554
                )
2555
            out["modules"][mod_name] = mod_name_short
×
2556
        else:
2557
            out["modules"][
1✔
2558
                "StarCatalog"
2559
            ] = self.TargetList.StarCatalog  # we just copy the StarCatalog string
2560

2561
        # if we don't know about the SurveyEnsemble, just write a blank to the output
2562
        if "SurveyEnsemble" not in out["modules"]:
1✔
2563
            out["modules"]["SurveyEnsemble"] = " "
1✔
2564

2565
        # get version and Git information
2566
        out["Version"] = get_version()
1✔
2567

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

2581
        return out
1✔
2582

2583
    def generateHashfName(self, specs):
1✔
2584
        """Generate cached file Hashname
2585

2586
        Requires a .XXX appended to end of hashname for each individual use case
2587

2588
        Args:
2589
            specs (dict):
2590
                :ref:`sec:inputspec`
2591

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

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

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

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

2652
    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
1✔
2653
        """Helper method for Overloading Revisit Filtering
2654

2655
        Args:
2656
            sInds (~numpy.ndarray(int)):
2657
                Indices of stars still in observation list
2658
            tmpCurrentTimeNorm (astropy.units.Quantity):
2659
                Normalized simulation time
2660

2661
        Returns:
2662
            ~numpy.ndarray(int):
2663
                indices of stars still in observation list
2664
        """
2665

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

2681
    def is_earthlike(self, plan_inds, sInd):
1✔
2682
        """Is the planet earthlike? Returns boolean array that's True for Earth-like
2683
        planets.
2684

2685

2686
        Args:
2687
            plan_inds(~numpy.ndarray(int)):
2688
                Planet indices
2689
            sInd (int):
2690
                Star index
2691

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

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

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

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

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

2739
        K = (
×
2740
            (c / np.sqrt(1 - e[t_filt]))
2741
            * Mpj[t_filt]
2742
            * np.sin(SU.I[t_filt])
2743
            * Ms[t_filt] ** (-2 / 3)
2744
            * T[t_filt] ** (-1 / 3)
2745
        )
2746

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

2752
        if PPop.scaleOrbits:
×
2753
            a_plan = (SU.a / np.sqrt(L_star)).value
×
2754
        else:
2755
            a_plan = SU.a.value
×
2756

2757
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
×
2758

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

2770
        # these are known_rv stars with earths around them
2771
        known_stars = np.unique(SU.plan2star[k_filt])  # these are known_rv stars
×
2772
        known_rocky = np.unique(SU.plan2star[r_filt])
×
2773

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

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

2784
        The observation window (which includes settling and overhead times)
2785
        is a superset of the integration window (in which photons are collected).
2786

2787
        The observation window begins at startTime. The integration window
2788
        begins at startTime + Obs.settlingTime + mode["ohTime"],
2789
        and the integration itself has a duration of intTime.
2790

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

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

2815
        """
2816

2817
        OS = self.OpticalSystem
×
2818
        ZL = self.ZodiacalLight
×
2819
        TL = self.TargetList
×
2820
        SU = self.SimulatedUniverse
×
2821
        Obs = self.Observatory
×
2822
        TK = self.TimeKeeping
×
2823

2824
        # time at start of integration window
2825
        currentTimeNorm = startTime
×
2826
        currentTimeAbs = TK.missionStart + startTime
×
2827

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

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

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

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

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

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

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

2917
        return SNR, fZ, systemParams
×
2918

2919
    def revisit_inds(self, sInds, tmpCurrentTimeNorm):
1✔
2920
        """Return subset of star indices that are scheduled for revisit.
2921

2922
        Args:
2923
            sInds (~numpy.ndarray(int)):
2924
                Indices of stars to consider
2925
            tmpCurrentTimeNorm (astropy.units.Quantity):
2926
                Normalized simulation time
2927

2928
        Returns:
2929
            ~numpy.ndarray(int):
2930
                Indices of stars whose revisit is scheduled for within `self.dt_max` of
2931
                the current time
2932

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

2939
        return ind_rev
1✔
2940

2941
    def keepout_filter(self, sInds, startTimes, endTimes, koMap):
1✔
2942
        """Filters stars not observable due to keepout
2943

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

2955
        Returns:
2956
            ~numpy.ndarray(int):
2957
                sInds - filtered indices of stars still in observation list
2958

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

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

2983
        return sInds[availableTargets]
×
2984

2985

2986
def array_encoder(obj):
1✔
2987
    r"""Encodes numpy arrays, astropy Times, and astropy Quantities, into JSON.
2988

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

2995
    Args:
2996
        obj (Any):
2997
            Object to encode.
2998

2999
        Returns:
3000
            Any:
3001
                Encoded object
3002

3003
    """
3004

3005
    from astropy.coordinates import SkyCoord
1✔
3006
    from astropy.time import Time
1✔
3007

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