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

dsavransky / EXOSIMS / 24730039167

21 Apr 2026 03:06PM UTC coverage: 66.015% (+0.2%) from 65.775%
24730039167

push

github

web-flow
Merge pull request #460 from dsavransky/tierScheduler_fixes

Tier scheduler fixes and new detection logics

6 of 6 new or added lines in 2 files covered. (100.0%)

14 existing lines in 4 files now uncovered.

9897 of 14992 relevant lines covered (66.02%)

0.66 hits per line

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

66.67
/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._numpy_compat import copy_if_needed
1✔
21
from EXOSIMS.util.deltaMag import deltaMag
1✔
22
from EXOSIMS.util.get_dirs import get_cache_dir
1✔
23
from EXOSIMS.util.get_module import get_module
1✔
24
from EXOSIMS.util.version_util import get_version
1✔
25
from EXOSIMS.util.vprint import vprint
1✔
26

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

29

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

33
    Args:
34

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

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

188
    """
189

190
    _modtype = "SurveySimulation"
1✔
191

192
    def __init__(
1✔
193
        self,
194
        scriptfile=None,
195
        ntFlux=1,
196
        nVisitsMax=5,
197
        charMargin=0.15,
198
        dt_max=1.0,
199
        record_counts_path=None,
200
        revisit_wait=None,
201
        nokoMap=False,
202
        nofZ=False,
203
        cachedir=None,
204
        defaultAddExoplanetObsTime=True,
205
        find_known_RV=False,
206
        include_known_RV=None,
207
        make_debug_bird_plots=False,
208
        debug_plot_path=None,
209
        **specs,
210
    ):
211
        # Commonly used units
212
        self.zero_d = 0 << u.d
1✔
213
        self.zero_arcsec = 0 * u.arcsec
1✔
214
        self.inv_arcsec2 = 1 / u.arcsec**2
1✔
215
        self.inv_s = 1 / u.s
1✔
216
        self.JEZ_unit = u.ph / u.s / u.m**2 / u.arcsec**2
1✔
217
        self.fZ_unit = 1 / u.arcsec**2
1✔
218
        self.AU2pc = (1 * u.AU).to_value(u.pc)
1✔
219
        self.rad2arcsec = (1 * u.rad).to_value(u.arcsec)
1✔
220
        self.day2sec = (1 * u.d).to_value(u.s)
1✔
221
        self.m_per_s = u.m / u.s
1✔
222

223
        # start the outspec
224
        self._outspec = {}
1✔
225

226
        # if a script file is provided read it in. If not set, assumes that
227
        # dictionary has been passed through specs.
228
        if scriptfile is not None:
1✔
229
            assert os.path.isfile(scriptfile), "%s is not a file." % scriptfile
1✔
230

231
            try:
1✔
232
                with open(scriptfile) as ff:
1✔
233
                    script = ff.read()
1✔
234
                specs.update(json.loads(script))
1✔
235
            except ValueError:
1✔
236
                sys.stderr.write("Script file `%s' is not valid JSON." % scriptfile)
1✔
237
                # must re-raise, or the error will be masked
238
                raise
1✔
239

240
            # modules array must be present
241
            if "modules" not in specs:
1✔
242
                raise ValueError("No modules field found in script.")
×
243

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

247
        # count dict contains all of the C info for each star index
248
        self.record_counts_path = record_counts_path
1✔
249
        self.count_lines = []
1✔
250
        self._outspec["record_counts_path"] = record_counts_path
1✔
251

252
        # mission simulation logger
253
        self.logger = specs.get("logger", logging.getLogger(__name__))
1✔
254

255
        # set up numpy random number (generate it if not in specs)
256
        self.seed = int(specs.get("seed", py_random.randint(1, int(1e9))))
1✔
257
        self.vprint("Numpy random seed is: %s" % self.seed)
1✔
258
        np.random.seed(self.seed)
1✔
259
        self._outspec["seed"] = self.seed
1✔
260

261
        # cache directory
262
        self.cachedir = get_cache_dir(cachedir)
1✔
263
        self._outspec["cachedir"] = self.cachedir
1✔
264
        # N.B.: cachedir is going to be used by everything, so let's make sure that
265
        # it doesn't get popped out of specs
266
        specs["cachedir"] = self.cachedir
1✔
267

268
        # if any of the modules is a string, assume that they are all strings
269
        # and we need to initalize
270
        if isinstance(next(iter(specs["modules"].values())), str):
1✔
271
            # import desired module names (prototype or specific)
272
            self.SimulatedUniverse = get_module(
1✔
273
                specs["modules"]["SimulatedUniverse"], "SimulatedUniverse"
274
            )(**specs)
275
            self.Observatory = get_module(
1✔
276
                specs["modules"]["Observatory"], "Observatory"
277
            )(**specs)
278
            self.TimeKeeping = get_module(
1✔
279
                specs["modules"]["TimeKeeping"], "TimeKeeping"
280
            )(**specs)
281

282
            # bring inherited class objects to top level of Survey Simulation
283
            SU = self.SimulatedUniverse
1✔
284
            self.StarCatalog = SU.StarCatalog
1✔
285
            self.PlanetPopulation = SU.PlanetPopulation
1✔
286
            self.PlanetPhysicalModel = SU.PlanetPhysicalModel
1✔
287
            self.OpticalSystem = SU.OpticalSystem
1✔
288
            self.ZodiacalLight = SU.ZodiacalLight
1✔
289
            self.BackgroundSources = SU.BackgroundSources
1✔
290
            self.PostProcessing = SU.PostProcessing
1✔
291
            self.Completeness = SU.Completeness
1✔
292
            self.TargetList = SU.TargetList
1✔
293

294
        else:
295
            # these are the modules that must be present if passing instantiated objects
296
            neededObjMods = [
1✔
297
                "PlanetPopulation",
298
                "PlanetPhysicalModel",
299
                "OpticalSystem",
300
                "ZodiacalLight",
301
                "BackgroundSources",
302
                "PostProcessing",
303
                "Completeness",
304
                "TargetList",
305
                "SimulatedUniverse",
306
                "Observatory",
307
                "TimeKeeping",
308
            ]
309

310
            # ensure that you have the minimal set
311
            for modName in neededObjMods:
1✔
312
                if modName not in specs["modules"]:
1✔
313
                    raise ValueError(
×
314
                        "%s module is required but was not provided." % modName
315
                    )
316

317
            for modName in specs["modules"]:
1✔
318
                assert specs["modules"][modName]._modtype == modName, (
1✔
319
                    "Provided instance of %s has incorrect modtype." % modName
320
                )
321

322
                setattr(self, modName, specs["modules"][modName])
1✔
323

324
        # create a dictionary of all modules, except StarCatalog
325
        self.modules = {}
1✔
326
        self.modules["PlanetPopulation"] = self.PlanetPopulation
1✔
327
        self.modules["PlanetPhysicalModel"] = self.PlanetPhysicalModel
1✔
328
        self.modules["OpticalSystem"] = self.OpticalSystem
1✔
329
        self.modules["ZodiacalLight"] = self.ZodiacalLight
1✔
330
        self.modules["BackgroundSources"] = self.BackgroundSources
1✔
331
        self.modules["PostProcessing"] = self.PostProcessing
1✔
332
        self.modules["Completeness"] = self.Completeness
1✔
333
        self.modules["TargetList"] = self.TargetList
1✔
334
        self.modules["SimulatedUniverse"] = self.SimulatedUniverse
1✔
335
        self.modules["Observatory"] = self.Observatory
1✔
336
        self.modules["TimeKeeping"] = self.TimeKeeping
1✔
337
        # add yourself to modules list for bookkeeping purposes
338
        self.modules["SurveySimulation"] = self
1✔
339

340
        # observation time sampling
341
        self.ntFlux = int(ntFlux)
1✔
342
        self._outspec["ntFlux"] = self.ntFlux
1✔
343

344
        # maximum number of observations per star
345
        self.nVisitsMax = int(nVisitsMax)
1✔
346
        self._outspec["nVisitsMax"] = self.nVisitsMax
1✔
347

348
        # integration time margin for characterization
349
        self.charMargin = float(charMargin)
1✔
350
        self._outspec["charMargin"] = self.charMargin
1✔
351

352
        # maximum time for revisit window
353
        self.dt_max = float(dt_max) * u.week
1✔
354
        self._outspec["dt_max"] = self.dt_max.value
1✔
355

356
        # minimum time for revisit window
357
        if revisit_wait is not None:
1✔
358
            self.revisit_wait = revisit_wait * u.d
×
359
        else:
360
            self.revisit_wait = revisit_wait
1✔
361
        self._outspec["revisit_wait"] = revisit_wait
1✔
362

363
        # list of detected earth-like planets aroung promoted stars
364
        self.known_earths = np.array([])
1✔
365

366
        self.find_known_RV = find_known_RV
1✔
367
        self._outspec["find_known_RV"] = find_known_RV
1✔
368
        self._outspec["include_known_RV"] = include_known_RV
1✔
369
        if self.find_known_RV:
1✔
370
            # select specific knonw RV stars if a file exists
371
            if include_known_RV is not None:
×
372
                if os.path.isfile(include_known_RV):
×
373
                    with open(include_known_RV, "r") as rv_file:
×
374
                        self.include_known_RV = [
×
375
                            hip.strip() for hip in rv_file.read().split(",")
376
                        ]
377
                        self.vprint(
×
378
                            "Including known RV stars: {}".format(self.include_known_RV)
379
                        )
380
                else:
381
                    self.include_known_RV = None
×
382
                    self.vprint(
×
383
                        "WARNING: Known RV file: {} does not exist!!".format(
384
                            include_known_RV
385
                        )
386
                    )
387
            else:
388
                self.include_known_RV = None
×
389
            self.known_stars, self.known_rocky = self.find_known_plans()
×
390
        else:
391
            self.include_known_RV = None
1✔
392
            self.known_stars = []
1✔
393
            self.known_rocky = []
1✔
394

395
        # defaultAddExoplanetObsTime Tells us time advanced when no targets available
396
        # counts agains exoplanetObsTime (when True)
397
        self.defaultAddExoplanetObsTime = defaultAddExoplanetObsTime
1✔
398
        self._outspec["defaultAddExoplanetObsTime"] = defaultAddExoplanetObsTime
1✔
399

400
        # If inputs are scalars, save scalars to outspec, otherwise save full lists
401
        OS = self.OpticalSystem
1✔
402
        TL = self.TargetList
1✔
403
        SU = self.SimulatedUniverse
1✔
404
        TK = self.TimeKeeping
1✔
405

406
        # initialize arrays updated in run_sim()
407
        self.initializeStorageArrays()
1✔
408

409
        # Generate File Hashnames and loction
410
        self.cachefname = self.generateHashfName(specs)
1✔
411

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

416
        # getting keepout map for entire mission
417
        startTime = self.TimeKeeping.missionStart.copy()
1✔
418
        endTime = self.TimeKeeping.missionFinishAbs.copy()
1✔
419

420
        nSystems = len(allModes)
1✔
421
        systNames = np.unique(
1✔
422
            [allModes[x]["syst"]["name"] for x in np.arange(nSystems)]
423
        )
424
        systOrder = np.argsort(systNames)
1✔
425
        koStr = ["koAngles_Sun", "koAngles_Moon", "koAngles_Earth", "koAngles_Small"]
1✔
426
        koangles = np.zeros([len(systNames), 4, 2])
1✔
427

428
        for x in systOrder:
1✔
429
            rel_mode = list(
1✔
430
                filter(lambda mode: mode["syst"]["name"] == systNames[x], allModes)
431
            )[0]
432
            koangles[x] = np.asarray([rel_mode["syst"][k] for k in koStr])
1✔
433

434
        # Precalculate Earth positions and create a cubic spline interpolator
435
        # to save time during orbit calculations
436
        dt = 0.1  # days
1✔
437
        self.Observatory.generate_Earth_interpolator(startTime, endTime, dt)
1✔
438

439
        self._outspec["nokoMap"] = nokoMap
1✔
440
        if not (nokoMap):
1✔
441
            koMaps, self.koTimes = self.Observatory.generate_koMap(
1✔
442
                TL, startTime, endTime, koangles
443
            )
444
            self.koTimes_mjd = self.koTimes.mjd
1✔
445
            self.koMaps = {}
1✔
446
            for x, n in zip(systOrder, systNames[systOrder]):
1✔
447
                print(n)
1✔
448
                self.koMaps[n] = koMaps[x, :, :]
1✔
449

450
        self._outspec["nofZ"] = nofZ
1✔
451

452
        # TODO: do we still want a nofZ option?  probably not.
453
        self.fZmins = {}
1✔
454
        self.fZtypes = {}
1✔
455
        for x, n in zip(systOrder, systNames[systOrder]):
1✔
456
            self.fZmins[n] = np.array([])
1✔
457
            self.fZtypes[n] = np.array([])
1✔
458

459
        # TODO: do we need to do this for all modes? doing det only breaks other
460
        # schedulers, but maybe there's a better approach here.
461
        sInds = np.arange(TL.nStars)  # Initialize some sInds array
1✔
462
        for mode in allModes:
1✔
463
            # This instantiates fZMap arrays for every starlight suppresion system
464
            # that is actually used in a mode
465
            modeHashName = (
1✔
466
                f'{self.cachefname[0:-1]}_{TL.nStars}_{mode["syst"]["name"]}'
467
                f"_{startTime.mjd:0}_{endTime.mjd:0}."
468
            )
469

470
            if (np.size(self.fZmins[mode["syst"]["name"]]) == 0) or (
1✔
471
                np.size(self.fZtypes[mode["syst"]["name"]]) == 0
472
            ):
473
                self.ZodiacalLight.generate_fZ(
1✔
474
                    self.Observatory,
475
                    TL,
476
                    self.TimeKeeping,
477
                    mode,
478
                    modeHashName,
479
                    self.koTimes,
480
                )
481

482
                (
1✔
483
                    self.fZmins[mode["syst"]["name"]],
484
                    self.fZtypes[mode["syst"]["name"]],
485
                ) = self.ZodiacalLight.calcfZmin(
486
                    sInds,
487
                    self.Observatory,
488
                    TL,
489
                    self.TimeKeeping,
490
                    mode,
491
                    modeHashName,
492
                    self.koMaps[mode["syst"]["name"]],
493
                    self.koTimes,
494
                )
495

496
        # Precalculating intTimeFilter for coronagraph
497
        # find fZmin to use in intTimeFilter
498
        self.valfZmin, self.absTimefZmin = self.ZodiacalLight.extractfZmin(
1✔
499
            self.fZmins[det_mode["syst"]["name"]], sInds, self.koTimes
500
        )
501
        JEZ0 = self.TargetList.JEZ0[det_mode["hex"]]
1✔
502
        dMag = TL.int_dMag[sInds]  # grabbing dMag
1✔
503
        WA = TL.int_WA[sInds]  # grabbing WA
1✔
504
        self.intTimesIntTimeFilter = (
1✔
505
            self.OpticalSystem.calc_intTime(
506
                TL, sInds, self.valfZmin, JEZ0, dMag, WA, det_mode, TK=TK
507
            )
508
            * det_mode["timeMultiplier"]
509
        )  # intTimes to filter by
510
        # These indices are acceptable for use simulating
511
        self.intTimeFilterInds = np.where(
1✔
512
            (
513
                (self.intTimesIntTimeFilter > 0)
514
                & (self.intTimesIntTimeFilter <= self.OpticalSystem.intCutoff)
515
            )
516
        )[0]
517

518
        self.make_debug_bird_plots = make_debug_bird_plots
1✔
519
        if self.make_debug_bird_plots:
1✔
520
            assert (
×
521
                debug_plot_path is not None
522
            ), "debug_plot_path must be set by input if make_debug_bird_plots is True"
523
            self.obs_plot_path = Path(f"{debug_plot_path}/{self.seed}")
×
524
            # Make directory if it doesn't exist
525
            if not self.obs_plot_path.exists():
×
526
                vprint(f"Making plot directory: {self.obs_plot_path}")
×
527
                self.obs_plot_path.mkdir(parents=True, exist_ok=True)
×
528
            self.obs_n_counter = 0
×
529
        self._outspec["make_debug_bird_plots"] = self.make_debug_bird_plots
1✔
530
        self._outspec["debug_plot_path"] = debug_plot_path
1✔
531

532
    def initializeStorageArrays(self):
1✔
533
        """
534
        Initialize all storage arrays based on # of stars and targets
535
        """
536

537
        self.DRM = []
1✔
538
        self.fullSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
1✔
539
        self.partialSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
1✔
540
        self.propagTimes = np.zeros(self.TargetList.nStars) * u.d
1✔
541
        self.lastObsTimes = np.zeros(self.TargetList.nStars) * u.d
1✔
542
        self.starVisits = np.zeros(
1✔
543
            self.TargetList.nStars, dtype=int
544
        )  # contains the number of times each star was visited
545
        self.starRevisit = np.array([])
1✔
546
        self.starExtended = np.array([], dtype=int)
1✔
547
        self.lastDetected = np.empty((self.TargetList.nStars, 4), dtype=object)
1✔
548

549
    def __str__(self):
1✔
550
        """String representation of the Survey Simulation object
551

552
        When the command 'print' is used on the Survey Simulation object, this
553
        method will return the values contained in the object
554

555
        """
556

557
        for att in self.__dict__:
1✔
558
            print("%s: %r" % (att, getattr(self, att)))
1✔
559

560
        return "Survey Simulation class object attributes"
1✔
561

562
    def run_sim(self):
1✔
563
        """Performs the survey simulation"""
564

565
        OS = self.OpticalSystem
1✔
566
        TL = self.TargetList
1✔
567
        SU = self.SimulatedUniverse
1✔
568
        Obs = self.Observatory
1✔
569
        TK = self.TimeKeeping
1✔
570

571
        # TODO: start using this self.currentSep
572
        # set occulter separation if haveOcculter
573
        if OS.haveOcculter:
1✔
574
            self.currentSep = Obs.occulterSep
×
575

576
        # choose observing modes selected for detection (default marked with a flag)
577
        allModes = OS.observingModes
1✔
578
        det_mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
579
        # and for characterization (default is first spectro/IFS mode)
580
        spectroModes = list(
1✔
581
            filter(lambda mode: "spec" in mode["inst"]["name"], allModes)
582
        )
583
        if np.any(spectroModes):
1✔
584
            char_mode = spectroModes[0]
×
585
        # if no spectro mode, default char mode is first observing mode
586
        else:
587
            char_mode = allModes[0]
1✔
588

589
        # begin Survey, and loop until mission is finished
590
        log_begin = "OB%s: survey beginning." % (TK.OBnumber)
1✔
591
        self.logger.info(log_begin)
1✔
592
        self.vprint(log_begin)
1✔
593
        t0 = time.time()
1✔
594
        sInd = None
1✔
595
        ObsNum = 0
1✔
596
        while not TK.mission_is_over(OS, Obs, det_mode):
1✔
597
            # acquire the NEXT TARGET star index and create DRM
598
            old_sInd = sInd  # used to save sInd if returned sInd is None
1✔
599
            DRM, sInd, det_intTime, waitTime = self.next_target(sInd, det_mode)
1✔
600

601
            if sInd is not None:
1✔
602
                ObsNum += (
1✔
603
                    1  # we're making an observation so increment observation number
604
                )
605

606
                if OS.haveOcculter:
1✔
607
                    # advance to start of observation (add slew time for selected target
608
                    _ = TK.advanceToAbsTime(TK.currentTimeAbs.copy() + waitTime)
×
609

610
                # beginning of observation, start to populate DRM
611
                DRM["star_ind"] = sInd
1✔
612
                DRM["star_name"] = TL.Name[sInd]
1✔
613
                DRM["arrival_time"] = TK.currentTimeNorm.to("day").copy()
1✔
614
                DRM["OB_nb"] = TK.OBnumber
1✔
615
                DRM["ObsNum"] = ObsNum
1✔
616
                pInds = np.where(SU.plan2star == sInd)[0]
1✔
617
                DRM["plan_inds"] = pInds.astype(int)
1✔
618
                log_obs = (
1✔
619
                    "  Observation #%s, star ind %s (of %s) with %s planet(s), "
620
                    + "mission time at Obs start: %s, exoplanetObsTime: %s"
621
                ) % (
622
                    ObsNum,
623
                    sInd,
624
                    TL.nStars,
625
                    len(pInds),
626
                    TK.currentTimeNorm.to("day").copy().round(2),
627
                    TK.exoplanetObsTime.to("day").copy().round(2),
628
                )
629
                self.logger.info(log_obs)
1✔
630
                self.vprint(log_obs)
1✔
631

632
                # PERFORM DETECTION and populate revisit list attribute
633
                (
1✔
634
                    detected,
635
                    det_fZ,
636
                    det_JEZ,
637
                    det_systemParams,
638
                    det_SNR,
639
                    FA,
640
                ) = self.observation_detection(sInd, det_intTime.copy(), det_mode)
641
                # update the occulter wet mass
642
                if OS.haveOcculter:
1✔
643
                    DRM = self.update_occulter_mass(
×
644
                        DRM, sInd, det_intTime.copy(), "det"
645
                    )
646
                # populate the DRM with detection results
647
                DRM["det_time"] = det_intTime.to("day")
1✔
648
                DRM["det_status"] = detected
1✔
649
                DRM["det_SNR"] = det_SNR
1✔
650
                DRM["det_fZ"] = det_fZ.to(self.fZ_unit)
1✔
651
                DRM["det_JEZ"] = det_JEZ.to(self.JEZ_unit)
1✔
652
                DRM["det_params"] = det_systemParams
1✔
653

654
                # PERFORM CHARACTERIZATION and populate spectra list attribute
655
                if char_mode["SNR"] not in [0, np.inf]:
1✔
656
                    (
1✔
657
                        characterized,
658
                        char_fZ,
659
                        char_JEZ,
660
                        char_systemParams,
661
                        char_SNR,
662
                        char_intTime,
663
                    ) = self.observation_characterization(sInd, char_mode)
664
                else:
665
                    char_intTime = None
×
666
                    lenChar = len(pInds) + 1 if FA else len(pInds)
×
667
                    characterized = np.zeros(lenChar, dtype=float)
×
668
                    char_SNR = np.zeros(lenChar, dtype=float)
×
669
                    char_fZ = 0.0 / u.arcsec**2
×
670
                    char_systemParams = SU.dump_system_params(sInd)
×
671
                assert char_intTime != 0, "Integration time can't be 0."
1✔
672
                # update the occulter wet mass
673
                if OS.haveOcculter and (char_intTime is not None):
1✔
674
                    DRM = self.update_occulter_mass(DRM, sInd, char_intTime, "char")
×
675
                # populate the DRM with characterization results
676
                DRM["char_time"] = (
1✔
677
                    char_intTime.to("day") if char_intTime is not None else 0.0 * u.day
678
                )
679
                DRM["char_status"] = characterized[:-1] if FA else characterized
1✔
680
                DRM["char_SNR"] = char_SNR[:-1] if FA else char_SNR
1✔
681
                DRM["char_fZ"] = char_fZ.to(self.fZ_unit)
1✔
682
                DRM["char_params"] = char_systemParams
1✔
683
                # populate the DRM with FA results
684
                DRM["FA_det_status"] = int(FA)
1✔
685
                DRM["FA_char_status"] = characterized[-1] if FA else 0
1✔
686
                DRM["FA_char_SNR"] = char_SNR[-1] if FA else 0.0
1✔
687
                DRM["FA_char_JEZ"] = (
1✔
688
                    self.lastDetected[sInd, 1][-1] if FA else 0.0 * self.JEZ_unit
689
                )
690
                DRM["FA_char_dMag"] = self.lastDetected[sInd, 2][-1] if FA else 0.0
1✔
691
                DRM["FA_char_WA"] = (
1✔
692
                    self.lastDetected[sInd, 3][-1] * u.arcsec if FA else 0.0 * u.arcsec
693
                )
694

695
                # populate the DRM with observation modes
696
                DRM["det_mode"] = dict(det_mode)
1✔
697
                del DRM["det_mode"]["inst"], DRM["det_mode"]["syst"]
1✔
698
                DRM["char_mode"] = dict(char_mode)
1✔
699
                del DRM["char_mode"]["inst"], DRM["char_mode"]["syst"]
1✔
700
                DRM["exoplanetObsTime"] = TK.exoplanetObsTime.copy()
1✔
701

702
                # append result values to self.DRM
703
                self.DRM.append(DRM)
1✔
704

705
                # handle case of inf OBs and missionPortion < 1
706
                if np.isinf(TK.OBduration) and (TK.missionPortion < 1.0):
1✔
707
                    self.arbitrary_time_advancement(
1✔
708
                        TK.currentTimeNorm.to("day").copy() - DRM["arrival_time"]
709
                    )
710

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

815
    def arbitrary_time_advancement(self, dt):
1✔
816
        """Handles fully dynamically scheduled case where OBduration is infinite and
817
        missionPortion is less than 1.
818

819
        Args:
820
            dt (~astropy.units.quantity.Quantity):
821
                Total amount of time, including all overheads and extras used for the
822
                previous observation.
823

824
        Returns:
825
            None
826
        """
827

828
        self.TimeKeeping.allocate_time(
1✔
829
            dt
830
            * (1.0 - self.TimeKeeping.missionPortion)
831
            / self.TimeKeeping.missionPortion,
832
            addExoplanetObsTime=False,
833
        )
834

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

838
        This method chooses the next target star index based on which
839
        stars are available, their integration time, and maximum completeness.
840
        Returns None if no target could be found.
841

842
        Args:
843
            old_sInd (int):
844
                Index of the previous target star
845
            mode (dict):
846
                Selected observing mode for detection
847

848
        Returns:
849
            tuple:
850
                DRM (dict):
851
                    Design Reference Mission, contains the results of one complete
852
                    observation (detection and characterization)
853
                sInd (int):
854
                    Index of next target star. Defaults to None.
855
                intTime (astropy.units.Quantity):
856
                    Selected star integration time for detection in units of day.
857
                    Defaults to None.
858
                waitTime (astropy.units.Quantity):
859
                    a strategically advantageous amount of time to wait in the case
860
                    of an occulter for slew times
861

862
        """
863
        OS = self.OpticalSystem
1✔
864
        TL = self.TargetList
1✔
865
        Obs = self.Observatory
1✔
866
        TK = self.TimeKeeping
1✔
867

868
        # create DRM
869
        DRM = {}
1✔
870

871
        # allocate settling time + overhead time
872
        tmpCurrentTimeAbs = (
1✔
873
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
874
        )
875
        tmpCurrentTimeNorm = (
1✔
876
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
877
        )
878

879
        # create appropriate koMap
880
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
881

882
        # look for available targets
883
        # 1. initialize arrays
884
        slewTimes = np.zeros(TL.nStars) << u.d
1✔
885
        # fZs = np.zeros(TL.nStars) / u.arcsec**2.0
886
        dV = np.zeros(TL.nStars) << self.m_per_s
1✔
887
        intTimes = np.zeros(TL.nStars) << u.d
1✔
888
        obsTimes = np.zeros([2, TL.nStars]) << u.d
1✔
889
        sInds = np.arange(TL.nStars)
1✔
890

891
        # 2. find spacecraft orbital START positions (if occulter, positions
892
        # differ for each star) and filter out unavailable targets
893
        sd = None
1✔
894
        if OS.haveOcculter:
1✔
895
            sd = Obs.star_angularSep(TL, old_sInd, sInds, tmpCurrentTimeAbs)
×
896
            obsTimes = Obs.calculate_observableTimes(
×
897
                TL, sInds, tmpCurrentTimeAbs, self.koMaps, self.koTimes, mode
898
            )
899
            slewTimes = Obs.calculate_slewTimes(
×
900
                TL, old_sInd, sInds, sd, obsTimes, tmpCurrentTimeAbs
901
            )
902

903
        # 2.1 filter out totTimes > integration cutoff
904
        if len(sInds.tolist()) > 0:
1✔
905
            sInds = np.intersect1d(self.intTimeFilterInds, sInds)
1✔
906

907
        # start times, including slew times
908
        startTimes = tmpCurrentTimeAbs.copy() + slewTimes
1✔
909
        startTimesNorm = tmpCurrentTimeNorm.copy() + slewTimes
1✔
910

911
        # 2.5 Filter stars not observable at startTimes
912
        try:
1✔
913
            tmpIndsbool = list()
1✔
914
            for i in np.arange(len(sInds)):
1✔
915
                koTimeInd = np.where(
1✔
916
                    np.round(startTimes[sInds[i]].value) - self.koTimes.value == 0
917
                )[0][
918
                    0
919
                ]  # find indice where koTime is startTime[0]
920
                tmpIndsbool.append(
1✔
921
                    koMap[sInds[i]][koTimeInd].astype(bool)
922
                )  # Is star observable at time ind
923
            sInds = sInds[tmpIndsbool]
1✔
924
            del tmpIndsbool
1✔
925
        except:  # noqa: E722 # If there are no target stars to observe
×
926
            sInds = np.asarray([], dtype=int)
×
927

928
        # 3. filter out all previously (more-)visited targets, unless in
929
        if len(sInds.tolist()) > 0:
1✔
930
            sInds = self.revisitFilter(sInds, tmpCurrentTimeNorm)
1✔
931

932
        # 4.1 calculate integration times for ALL preselected targets
933
        (
1✔
934
            maxIntTimeOBendTime,
935
            maxIntTimeExoplanetObsTime,
936
            maxIntTimeMissionLife,
937
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
938
        maxIntTime = min(
1✔
939
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
940
        )  # Maximum intTime allowed
941

942
        if len(sInds.tolist()) > 0:
1✔
943
            if OS.haveOcculter and (old_sInd is not None):
1✔
944
                (
×
945
                    sInds,
946
                    slewTimes[sInds],
947
                    intTimes[sInds],
948
                    dV[sInds],
949
                ) = self.refineOcculterSlews(
950
                    old_sInd, sInds, slewTimes, obsTimes, sd, mode
951
                )
952
                endTimes = tmpCurrentTimeAbs.copy() + intTimes + slewTimes
×
953
            else:
954
                intTimes[sInds] = self.calc_targ_intTime(sInds, startTimes[sInds], mode)
1✔
955
                sInds = sInds[
1✔
956
                    np.where(intTimes[sInds] <= maxIntTime)
957
                ]  # Filters targets exceeding end of OB
958
                endTimes = tmpCurrentTimeAbs.copy() + intTimes
1✔
959

960
                if maxIntTime.value <= 0:
1✔
961
                    sInds = np.asarray([], dtype=int)
×
962

963
        # 5.1 TODO Add filter to filter out stars entering and exiting keepout
964
        # between startTimes and endTimes
965

966
        # 5.2 find spacecraft orbital END positions (for each candidate target),
967
        # and filter out unavailable targets
968
        if len(sInds.tolist()) > 0 and Obs.checkKeepoutEnd:
1✔
969
            # endTimes may exist past koTimes so we have an exception to hand this case
970
            try:
1✔
971
                tmpIndsbool = list()
1✔
972
                for i in np.arange(len(sInds)):
1✔
973
                    koTimeInd = np.where(
1✔
974
                        np.round(endTimes[sInds[i]].value) - self.koTimes.value == 0
975
                    )[0][
976
                        0
977
                    ]  # find indice where koTime is endTime[0]
978
                    tmpIndsbool.append(
1✔
979
                        koMap[sInds[i]][koTimeInd].astype(bool)
980
                    )  # Is star observable at time ind
981
                sInds = sInds[tmpIndsbool]
1✔
982
                del tmpIndsbool
1✔
983
            except:  # noqa: E722
×
984
                sInds = np.asarray([], dtype=int)
×
985

986
        # 6. choose best target from remaining
987
        if len(sInds.tolist()) > 0:
1✔
988
            # choose sInd of next target
989
            sInd, waitTime = self.choose_next_target(
1✔
990
                old_sInd, sInds, slewTimes, intTimes[sInds]
991
            )
992

993
            # Should Choose Next Target decide there are no stars it wishes to
994
            # observe at this time.
995
            if sInd is None and (waitTime is not None):
1✔
996
                self.vprint(
×
997
                    "There are no stars available to observe. Waiting {}".format(
998
                        waitTime
999
                    )
1000
                )
1001
                return DRM, None, None, waitTime
×
1002
            elif (sInd is None) and (waitTime is None):
1✔
1003
                self.vprint(
×
1004
                    "There are no stars available to observe and waitTime is None."
1005
                )
1006
                return DRM, None, None, waitTime
×
1007
            # store selected star integration time
1008
            intTime = intTimes[sInd]
1✔
1009

1010
        # if no observable target, advanceTime to next Observable Target
1011
        else:
1012
            self.vprint(
×
1013
                "No Observable Targets at currentTimeNorm= "
1014
                + str(TK.currentTimeNorm.copy())
1015
            )
1016
            return DRM, None, None, None
×
1017

1018
        # update visited list for selected star
1019
        self.starVisits[sInd] += 1
1✔
1020
        # store normalized start time for future completeness update
1021
        self.lastObsTimes[sInd] = startTimesNorm[sInd]
1✔
1022

1023
        # populate DRM with occulter related values
1024
        if OS.haveOcculter:
1✔
1025
            DRM = Obs.log_occulterResults(
×
1026
                DRM, slewTimes[sInd], sInd, sd[sInd], dV[sInd]
1027
            )
1028
            return DRM, sInd, intTime, slewTimes[sInd]
×
1029

1030
        return DRM, sInd, intTime, waitTime
1✔
1031

1032
    def calc_targ_intTime(self, sInds, startTimes, mode):
1✔
1033
        """Helper method for next_target to aid in overloading for alternative
1034
        implementations.
1035

1036
        Given a subset of targets, calculate their integration times given the
1037
        start of observation time.
1038

1039
        Prototype just calculates integration times for fixed contrast depth.
1040

1041
        Args:
1042
            sInds (~numpy.ndarray(int)):
1043
                Indices of available targets
1044
            startTimes (astropy quantity array):
1045
                absolute start times of observations.
1046
                must be of the same size as sInds
1047
            mode (dict):
1048
                Selected observing mode for detection
1049

1050
        Returns:
1051
            ~astropy.units.Quantity(~numpy.ndarray(float)):
1052
                Integration times for detection
1053
                same dimension as sInds
1054

1055
        .. note::
1056

1057
            next_target filter will discard targets with zero integration times.
1058

1059
        """
1060

1061
        TL = self.TargetList
1✔
1062

1063
        # assumed values for detection
1064
        fZ = self.ZodiacalLight.fZ(
1✔
1065
            self.Observatory, self.TargetList, sInds, startTimes, mode
1066
        )
1067
        # Use the default exozodi intensities
1068
        JEZ = TL.JEZ0[mode["hex"]][sInds]
1✔
1069
        dMag = TL.int_dMag[sInds]
1✔
1070
        WA = TL.int_WA[sInds]
1✔
1071

1072
        # save out file containing photon count info
1073
        if self.record_counts_path is not None and len(self.count_lines) == 0:
1✔
1074
            C_p, C_b, C_sp, C_extra = self.OpticalSystem.Cp_Cb_Csp(
×
1075
                self.TargetList, sInds, fZ, JEZ, dMag, WA, mode, returnExtra=True
1076
            )
1077
            import csv
×
1078

1079
            count_fpath = os.path.join(self.record_counts_path, "counts")
×
1080

1081
            if not os.path.exists(count_fpath):
×
1082
                os.mkdir(count_fpath)
×
1083

1084
            outfile = os.path.join(count_fpath, str(self.seed) + ".csv")
×
1085
            self.count_lines.append(
×
1086
                [
1087
                    "sInd",
1088
                    "HIPs",
1089
                    "C_F0",
1090
                    "C_p0",
1091
                    "C_sr",
1092
                    "C_z",
1093
                    "C_ez",
1094
                    "C_dc",
1095
                    "C_cc",
1096
                    "C_rn",
1097
                    "C_p",
1098
                    "C_b",
1099
                    "C_sp",
1100
                ]
1101
            )
1102

1103
            for i, sInd in enumerate(sInds):
×
1104
                self.count_lines.append(
×
1105
                    [
1106
                        sInd,
1107
                        self.TargetList.Name[sInd],
1108
                        C_extra["C_F0"][0].value,
1109
                        C_extra["C_sr"][i].value,
1110
                        C_extra["C_z"][i].value,
1111
                        C_extra["C_ez"][i].value,
1112
                        C_extra["C_dc"][i].value,
1113
                        C_extra["C_cc"][i].value,
1114
                        C_extra["C_rn"][i].value,
1115
                        C_p[i].value,
1116
                        C_b[i].value,
1117
                        C_sp[i].value,
1118
                    ]
1119
                )
1120

1121
            with open(outfile, "w") as csvfile:
×
1122
                c = csv.writer(csvfile)
×
1123
                c.writerows(self.count_lines)
×
1124

1125
        intTimes = self.OpticalSystem.calc_intTime(
1✔
1126
            self.TargetList, sInds, fZ, JEZ, dMag, WA, mode
1127
        )
1128
        intTimes[~np.isfinite(intTimes)] = 0 * u.d
1✔
1129

1130
        return intTimes
1✔
1131

1132
    def choose_next_target(self, old_sInd, sInds, slewTimes, intTimes):
1✔
1133
        """Helper method for method next_target to simplify alternative implementations.
1134

1135
        Given a subset of targets (pre-filtered by method next_target or some
1136
        other means), select the best next one. The prototype uses completeness
1137
        as the sole heuristic.
1138

1139
        Args:
1140
            old_sInd (int):
1141
                Index of the previous target star
1142
            sInds (~numpy.ndarray(int)):
1143
                Indices of available targets
1144
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1145
                slew times to all stars (must be indexed by sInds)
1146
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1147
                Integration times for detection in units of day
1148

1149
        Returns:
1150
            tuple:
1151
                sInd (int):
1152
                    Index of next target star
1153
                waitTime (:py:class:`~astropy.units.Quantity`):
1154
                    Some strategic amount of time to wait in case an occulter slew is
1155
                    desired (default is None)
1156

1157
        """
1158

1159
        Comp = self.Completeness
1✔
1160
        TL = self.TargetList
1✔
1161
        TK = self.TimeKeeping
1✔
1162
        OS = self.OpticalSystem
1✔
1163
        Obs = self.Observatory
1✔
1164
        allModes = OS.observingModes
1✔
1165

1166
        # cast sInds to array
1167
        sInds = np.array(sInds, ndmin=1, copy=copy_if_needed)
1✔
1168
        # calculate dt since previous observation
1169
        dt = TK.currentTimeNorm.copy() + slewTimes[sInds] - self.lastObsTimes[sInds]
1✔
1170
        # get dynamic completeness values
1171
        comps = Comp.completeness_update(TL, sInds, self.starVisits[sInds], dt)
1✔
1172
        # choose target with maximum completeness
1173
        sInd = np.random.choice(sInds[comps == max(comps)])
1✔
1174

1175
        # Check if exoplanetObsTime would be exceeded
1176
        mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
1177
        (
1✔
1178
            maxIntTimeOBendTime,
1179
            maxIntTimeExoplanetObsTime,
1180
            maxIntTimeMissionLife,
1181
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
1182
        maxIntTime = min(
1✔
1183
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1184
        )  # Maximum intTime allowed
1185
        intTimes2 = self.calc_targ_intTime(
1✔
1186
            np.array([sInd]), TK.currentTimeAbs.copy(), mode
1187
        )
1188
        if (
1✔
1189
            intTimes2 > maxIntTime
1190
        ):  # check if max allowed integration time would be exceeded
1191
            self.vprint("max allowed integration time would be exceeded")
×
1192
            sInd = None
×
1193
            waitTime = 1.0 * u.d
×
1194
            return sInd, waitTime
×
1195

1196
        return (
1✔
1197
            sInd,
1198
            slewTimes[sInd],
1199
        )  # if coronagraph or first sInd, waitTime will be 0 days
1200

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

1204
        Refines the selection of occulter slew times by filtering based on mission time
1205
        constraints and selecting the best slew time for each star. This method calls on
1206
        other occulter methods within SurveySimulation depending on how slew times were
1207
        calculated prior to calling this function (i.e. depending on which
1208
        implementation of the Observatory module is used).
1209

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

1228
        Returns:
1229
            tuple:
1230
                sInds (int):
1231
                    Indeces of next target star
1232
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1233
                    slew times to all stars (must be indexed by sInds)
1234
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1235
                    Integration times for detection in units of day
1236
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1237
                    Delta-V used to transfer to new star line of sight in unis of m/s
1238
        """
1239

1240
        Obs = self.Observatory
×
1241
        TL = self.TargetList
×
1242

1243
        # initializing arrays
1244
        obsTimeArray = np.zeros([TL.nStars, 50]) << u.d
×
1245
        intTimeArray = np.zeros([TL.nStars, 2]) << u.d
×
1246

1247
        for n in sInds:
×
1248
            obsTimeArray[n, :] = (
×
1249
                np.linspace(obsTimes[0, n].value, obsTimes[1, n].value, 50) << u.d
1250
            )
1251
        intTimeArray[sInds, 0] = self.calc_targ_intTime(
×
1252
            sInds, Time(obsTimeArray[sInds, 0], format="mjd", scale="tai"), mode
1253
        )
1254
        intTimeArray[sInds, 1] = self.calc_targ_intTime(
×
1255
            sInds, Time(obsTimeArray[sInds, -1], format="mjd", scale="tai"), mode
1256
        )
1257

1258
        # determining which scheme to use to filter slews
1259
        obsModName = Obs.__class__.__name__
×
1260

1261
        # slew times have not been calculated/decided yet (SotoStarshade)
1262
        if obsModName == "SotoStarshade":
×
1263
            sInds, intTimes, slewTimes, dV = self.findAllowableOcculterSlews(
×
1264
                sInds,
1265
                old_sInd,
1266
                sd[sInds],
1267
                slewTimes[sInds],
1268
                obsTimeArray[sInds, :],
1269
                intTimeArray[sInds, :],
1270
                mode,
1271
            )
1272

1273
        # slew times were calculated/decided beforehand (Observatory Prototype)
1274
        else:
1275
            sInds, intTimes, slewTimes = self.filterOcculterSlews(
×
1276
                sInds,
1277
                slewTimes[sInds],
1278
                obsTimeArray[sInds, :],
1279
                intTimeArray[sInds, :],
1280
                mode,
1281
            )
1282
            dV = np.zeros(len(sInds)) << self.m_per_s
×
1283

1284
        return sInds, slewTimes, intTimes, dV
×
1285

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

1289
        Used by the refineOcculterSlews method when slew times have been selected
1290
        a priori. This method filters out slews that are not within desired observing
1291
        blocks, the maximum allowed integration time, and are outside of future
1292
        keepouts.
1293

1294
        Args:
1295
            sInds (~numpy.ndarray(int)):
1296
                Indices of available targets
1297
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1298
                slew times to all stars (must be indexed by sInds)
1299
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1300
                Array of times during which a star is out of keepout, has shape
1301
                nx50 where n is the number of stars in sInds. Unit of days
1302
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1303
                Array of integration times for each time in obsTimeArray, has shape
1304
                nx2 where n is the number of stars in sInds. Unit of days
1305
            mode (dict):
1306
                Selected observing mode for detection
1307

1308
        Returns:
1309
            tuple:
1310
                sInds (int):
1311
                    Indeces of next target star
1312
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1313
                    Integration times for detection in units of day
1314
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1315
                    slew times to all stars (must be indexed by sInds)
1316
        """
1317

1318
        TK = self.TimeKeeping
×
1319
        Obs = self.Observatory
×
1320

1321
        # allocate settling time + overhead time
1322
        tmpCurrentTimeAbs = (
×
1323
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1324
        )
1325
        tmpCurrentTimeNorm = (
×
1326
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1327
        )
1328

1329
        # 0. lambda function that linearly interpolates Integration Time
1330
        # between obsTimes
1331
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1332
            t - np.array(x[:, 0]).reshape(len(t), 1)
1333
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1334

1335
        # 1. initializing arrays
1336
        obsTimesRange = np.array(
×
1337
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1338
        )  # nx2 array with start and end times of obsTimes for each star
1339
        intTimesRange = np.array([intTimeArray[:, 0], intTimeArray[:, -1]])
×
1340

1341
        OBnumbers = np.zeros(
×
1342
            [len(sInds), 1]
1343
        )  # for each sInd, will show during which OB observations will take place
1344
        maxIntTimes = np.zeros([len(sInds), 1]) << u.d
×
1345

1346
        intTimes = (
×
1347
            linearInterp(
1348
                intTimesRange.T,
1349
                obsTimesRange.T,
1350
                (tmpCurrentTimeAbs + slewTimes).reshape(len(sInds), 1).value,
1351
            )
1352
            << u.d
1353
        )  # calculate intTimes for each slew time
1354

1355
        minObsTimeNorm = (obsTimesRange[0, :] - tmpCurrentTimeAbs.value).reshape(
×
1356
            [len(sInds), 1]
1357
        )
1358
        maxObsTimeNorm = (obsTimesRange[1, :] - tmpCurrentTimeAbs.value).reshape(
×
1359
            [len(sInds), 1]
1360
        )
1361
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1362

1363
        # 2. find OBnumber for each sInd's slew time
1364
        if len(TK.OBendTimes) > 1:
×
1365
            for i in range(len(sInds)):
×
1366
                S = np.where(
×
1367
                    TK.OBstartTimes.value - tmpCurrentTimeNorm.value
1368
                    < slewTimes[i].value
1369
                )[0][-1]
1370
                F = np.where(
×
1371
                    TK.OBendTimes.value - tmpCurrentTimeNorm.value < slewTimes[i].value
1372
                )[0]
1373

1374
                # case when slews are in the first OB
1375
                if F.shape[0] == 0:
×
1376
                    F = -1
×
1377
                else:
1378
                    F = F[-1]
×
1379

1380
                # slew occurs within an OB (nth OB has started but hasn't ended)
1381
                if S != F:
×
1382
                    OBnumbers[i] = S
×
1383
                    (
×
1384
                        maxIntTimeOBendTime,
1385
                        maxIntTimeExoplanetObsTime,
1386
                        maxIntTimeMissionLife,
1387
                    ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, TK.OBstartTimes[S], S)
1388
                    maxIntTimes[i] = min(
×
1389
                        maxIntTimeOBendTime,
1390
                        maxIntTimeExoplanetObsTime,
1391
                        maxIntTimeMissionLife,
1392
                    )  # Maximum intTime allowed
1393

1394
                # slew occurs between OBs, badbadnotgood
1395
                else:
1396
                    OBnumbers[i] = -1
×
1397
                    maxIntTimes[i] = 0 * u.d
×
1398
            OBstartTimeNorm = (
×
1399
                TK.OBstartTimes[np.array(OBnumbers, dtype=int)].value
1400
                - tmpCurrentTimeNorm.value
1401
            )
1402
        else:
1403
            (
×
1404
                maxIntTimeOBendTime,
1405
                maxIntTimeExoplanetObsTime,
1406
                maxIntTimeMissionLife,
1407
            ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm)
1408
            maxIntTimes[:] = min(
×
1409
                maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1410
            )  # Maximum intTime allowed
1411
            OBstartTimeNorm = np.zeros(OBnumbers.shape)
×
1412

1413
        # finding the minimum possible slew time
1414
        # (either OBstart or when star JUST leaves keepout)
1415
        minAllowedSlewTime = np.max([OBstartTimeNorm, minObsTimeNorm], axis=0)
×
1416

1417
        # 3. start filtering process
1418
        good_inds = np.where((OBnumbers >= 0) & (ObsTimeRange > intTimes.value))[0]
×
1419
        # star slew within OB -AND- can finish observing
1420
        # before it goes back into keepout
1421
        if good_inds.shape[0] > 0:
×
1422
            # the good ones
1423
            sInds = sInds[good_inds]
×
1424
            slewTimes = slewTimes[good_inds]
×
1425
            intTimes = intTimes[good_inds]
×
1426
            OBstartTimeNorm = OBstartTimeNorm[good_inds]
×
1427
            minAllowedSlewTime = minAllowedSlewTime[good_inds]
×
1428

1429
            # maximum allowed slew time based on integration times
1430
            maxAllowedSlewTime = maxIntTimes[good_inds].value - intTimes.value
×
1431
            maxAllowedSlewTime[maxAllowedSlewTime < 0] = -np.inf
×
1432
            maxAllowedSlewTime += OBstartTimeNorm  # calculated rel to currentTime norm
×
1433

1434
            # checking to see if slewTimes are allowed
1435
            good_inds = np.where(
×
1436
                (slewTimes.reshape([len(sInds), 1]).value > minAllowedSlewTime)
1437
                & (slewTimes.reshape([len(sInds), 1]).value < maxAllowedSlewTime)
1438
            )[0]
1439

1440
            slewTimes = slewTimes[good_inds]
×
1441
        else:
1442
            slewTimes = slewTimes[good_inds]
×
1443

1444
        return sInds[good_inds], intTimes[good_inds].flatten(), slewTimes
×
1445

1446
    def findAllowableOcculterSlews(
1✔
1447
        self, sInds, old_sInd, sd, slewTimes, obsTimeArray, intTimeArray, mode
1448
    ):
1449
        """Finds an array of allowable slew times for each star
1450

1451
        Used by the refineOcculterSlews method when slew times have NOT been selected
1452
        a priori. This method creates nx50 arrays (where the row corresponds to a
1453
        specific star and the column corresponds to a future point in time relative to
1454
        currentTime).
1455

1456
        These arrays are initially zero but are populated with the corresponding values
1457
        (slews, intTimes, etc) if slewing to that time point (i.e. beginning an
1458
        observation) would lead to a successful observation. A "successful observation"
1459
        is defined by certain conditions relating to keepout and the komap, observing
1460
        blocks, mission lifetime, and some constraints on the dVmap calculation in
1461
        SotoStarshade. Each star will likely have a range of slewTimes that would lead
1462
        to a successful observation -- another method is then called to select the best
1463
        of these slewTimes.
1464

1465
        Args:
1466
            sInds (~numpy.ndarray(int)):
1467
                Indices of available targets
1468
            old_sInd (int):
1469
                Index of the previous target star
1470
            sd (~astropy.units.Quantity(~numpy.ndarray(float))):
1471
                Angular separation between stars in rad
1472
            slewTimes (astropy quantity array):
1473
                slew times to all stars (must be indexed by sInds)
1474
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1475
                Array of times during which a star is out of keepout, has shape
1476
                nx50 where n is the number of stars in sInds
1477
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1478
                Array of integration times for each time in obsTimeArray, has shape
1479
                nx50 where n is the number of stars in sInds
1480
            mode (dict):
1481
                Selected observing mode for detection
1482

1483
        Returns:
1484
            tuple:
1485
                sInds (numpy.ndarray(int)):
1486
                    Indices of next target star
1487
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1488
                    slew times to all stars (must be indexed by sInds)
1489
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1490
                    Integration times for detection in units of day
1491
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1492
                    Delta-V used to transfer to new star line of sight in unis of m/s
1493
        """
1494
        TK = self.TimeKeeping
×
1495
        Obs = self.Observatory
×
1496
        TL = self.TargetList
×
1497

1498
        # 0. lambda function that linearly interpolates Integration
1499
        # Time between obsTimes
1500
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1501
            t - np.array(x[:, 0]).reshape(len(t), 1)
1502
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1503

1504
        # allocate settling time + overhead time
1505
        tmpCurrentTimeAbs = (
×
1506
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1507
        )
1508
        tmpCurrentTimeNorm = (
×
1509
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1510
        )
1511

1512
        # 1. initializing arrays
1513
        obsTimes = np.array(
×
1514
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1515
        )  # nx2 array with start and end times of obsTimes for each star
1516
        intTimes_int = (
×
1517
            np.zeros(obsTimeArray.shape) << u.d
1518
        )  # initializing intTimes of shape nx50 then interpolating
1519
        intTimes_int = (
×
1520
            np.hstack(
1521
                [
1522
                    intTimeArray[:, 0].reshape(len(sInds), 1).value,
1523
                    linearInterp(
1524
                        intTimeArray.value, obsTimes.T, obsTimeArray[:, 1:].value
1525
                    ),
1526
                ]
1527
            )
1528
            << u.d
1529
        )
1530
        allowedSlewTimes = np.zeros(obsTimeArray.shape) << u.d
×
1531
        allowedintTimes = np.zeros(obsTimeArray.shape) << u.d
×
1532
        allowedCharTimes = np.zeros(obsTimeArray.shape) << u.d
×
1533
        obsTimeArrayNorm = obsTimeArray.value - tmpCurrentTimeAbs.value
×
1534

1535
        # obsTimes -> relative to current Time
1536
        try:
×
1537
            minObsTimeNorm = np.array([np.min(v[v > 0]) for v in obsTimeArrayNorm])
×
1538
        except:  # noqa: E722
×
1539
            # an error pops up sometimes at the end of the mission, this fixes it
1540
            # TODO: define the error type that occurs
1541
            # rewrite to avoid a try/except if possible
1542
            minObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1543

1544
        maxObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1545
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1546

1547
        # getting max possible intTime
1548
        (
×
1549
            maxIntTimeOBendTime,
1550
            maxIntTimeExoplanetObsTime,
1551
            maxIntTimeMissionLife,
1552
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm, TK.OBnumber)
1553
        maxIntTime = min(
×
1554
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1555
        )  # Maximum intTime allowed
1556

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

1561
        # first filled in for the current OB
1562
        minAllowedSlewTimes = np.array(
×
1563
            [minObsTimeNorm.T] * len(intTimes_int.T)
1564
        ).T  # just to make it nx50 so it plays nice with the other arrays
1565
        maxAllowedSlewTimes = maxIntTime.value - intTimes_int.value
×
1566
        maxAllowedSlewTimes[maxAllowedSlewTimes > Obs.occ_dtmax.value] = (
×
1567
            Obs.occ_dtmax.value
1568
        )
1569

1570
        # conditions that must be met to define an allowable slew time
1571
        cond1 = (
×
1572
            minAllowedSlewTimes >= Obs.occ_dtmin.value
1573
        )  # minimum dt time in dV map interpolant
1574
        cond2 = (
×
1575
            maxAllowedSlewTimes <= Obs.occ_dtmax.value
1576
        )  # maximum dt time in dV map interpolant
1577
        cond3 = maxAllowedSlewTimes > minAllowedSlewTimes
×
1578
        cond4 = intTimes_int.value < ObsTimeRange.reshape(len(sInds), 1)
×
1579

1580
        conds = cond1 & cond2 & cond3 & cond4
×
1581
        minAllowedSlewTimes[np.invert(conds)] = (
×
1582
            np.inf
1583
        )  # these are filtered during the next filter
1584
        maxAllowedSlewTimes[np.invert(conds)] = -np.inf
×
1585

1586
        # one last condition to meet
1587
        map_i, map_j = np.where(
×
1588
            (obsTimeArrayNorm > minAllowedSlewTimes)
1589
            & (obsTimeArrayNorm < maxAllowedSlewTimes)
1590
        )
1591

1592
        # 2.5 if any stars are slew-able to within this OB block, populate
1593
        # "allowedSlewTimes", a running tally of possible slews
1594
        # within the time range a star is observable (out of keepout)
1595
        if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1596
            allowedSlewTimes[map_i, map_j] = obsTimeArrayNorm[map_i, map_j] * u.d
×
1597
            allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1598
            allowedCharTimes[map_i, map_j] = maxIntTime - intTimes_int[map_i, map_j]
×
1599

1600
        # 3. search future OBs
1601
        OB_withObsStars = (
×
1602
            TK.OBstartTimes.value - np.min(obsTimeArrayNorm) - tmpCurrentTimeNorm.value
1603
        )  # OBs within which any star is observable
1604

1605
        if any(OB_withObsStars > 0):
×
1606
            nOBstart = np.argmin(np.abs(OB_withObsStars))
×
1607
            nOBend = np.argmax(OB_withObsStars)
×
1608

1609
            # loop through the next 5 OBs (or until mission is over if there are less
1610
            # than 5 OBs in the future)
1611
            for i in np.arange(nOBstart, np.min([nOBend, nOBstart + 5])):
×
1612
                # max int Times for the next OB
1613
                (
×
1614
                    maxIntTimeOBendTime,
1615
                    maxIntTimeExoplanetObsTime,
1616
                    maxIntTimeMissionLife,
1617
                ) = TK.get_ObsDetectionMaxIntTime(
1618
                    Obs, mode, TK.OBstartTimes[i + 1], i + 1
1619
                )
1620
                maxIntTime_nOB = min(
×
1621
                    maxIntTimeOBendTime,
1622
                    maxIntTimeExoplanetObsTime,
1623
                    maxIntTimeMissionLife,
1624
                )  # Maximum intTime allowed
1625

1626
                # min and max slew times rel to current Time (norm)
1627
                nOBstartTimeNorm = np.array(
×
1628
                    [TK.OBstartTimes[i + 1].value - tmpCurrentTimeNorm.value]
1629
                    * len(sInds)
1630
                )
1631

1632
                # min slew times for stars start either whenever the star first leaves
1633
                # keepout or when next OB stars, whichever comes last
1634
                minAllowedSlewTimes_nOB = np.array(
×
1635
                    [np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T]
1636
                    * len(maxAllowedSlewTimes.T)
1637
                ).T
1638
                maxAllowedSlewTimes_nOB = (
×
1639
                    nOBstartTimeNorm.reshape(len(sInds), 1)
1640
                    + maxIntTime_nOB.value
1641
                    - intTimes_int.value
1642
                )
1643
                maxAllowedSlewTimes_nOB[
×
1644
                    maxAllowedSlewTimes_nOB > Obs.occ_dtmax.value
1645
                ] = Obs.occ_dtmax.value
1646

1647
                # amount of time left when the stars are still out of keepout
1648
                ObsTimeRange_nOB = (
×
1649
                    maxObsTimeNorm
1650
                    - np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T
1651
                )
1652

1653
                # condition to be met for an allowable slew time
1654
                cond1 = minAllowedSlewTimes_nOB >= Obs.occ_dtmin.value
×
1655
                cond2 = maxAllowedSlewTimes_nOB <= Obs.occ_dtmax.value
×
1656
                cond3 = maxAllowedSlewTimes_nOB > minAllowedSlewTimes_nOB
×
1657
                cond4 = intTimes_int.value < ObsTimeRange_nOB.reshape(len(sInds), 1)
×
1658
                cond5 = intTimes_int.value < maxIntTime_nOB.value
×
1659
                conds = cond1 & cond2 & cond3 & cond4 & cond5
×
1660

1661
                minAllowedSlewTimes_nOB[np.invert(conds)] = np.inf
×
1662
                maxAllowedSlewTimes_nOB[np.invert(conds)] = -np.inf
×
1663

1664
                # one last condition
1665
                map_i, map_j = np.where(
×
1666
                    (obsTimeArrayNorm > minAllowedSlewTimes_nOB)
1667
                    & (obsTimeArrayNorm < maxAllowedSlewTimes_nOB)
1668
                )
1669

1670
                # 3.33 populate the running tally of allowable slew times if it meets
1671
                # all conditions
1672
                if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1673
                    allowedSlewTimes[map_i, map_j] = (
×
1674
                        obsTimeArrayNorm[map_i, map_j] * u.d
1675
                    )
1676
                    allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1677
                    allowedCharTimes[map_i, map_j] = (
×
1678
                        maxIntTime_nOB - intTimes_int[map_i, map_j]
1679
                    )
1680

1681
        # 3.67 filter out any stars that are not observable at all
1682
        filterDuds = np.sum(allowedSlewTimes, axis=1) > 0.0
×
1683
        sInds = sInds[filterDuds]
×
1684

1685
        # 4. choose a slew time for each available star
1686
        # calculate dVs for each possible slew time for each star
1687
        allowed_dVs = Obs.calculate_dV(
×
1688
            TL,
1689
            old_sInd,
1690
            sInds,
1691
            sd[filterDuds],
1692
            allowedSlewTimes[filterDuds, :],
1693
            tmpCurrentTimeAbs,
1694
        )
1695

1696
        if len(sInds.tolist()) > 0:
×
1697
            # select slew time for each star
1698
            dV_inds = np.arange(0, len(sInds))
×
1699
            sInds, intTime, slewTime, dV = self.chooseOcculterSlewTimes(
×
1700
                sInds,
1701
                allowedSlewTimes[filterDuds, :],
1702
                allowed_dVs[dV_inds, :],
1703
                allowedintTimes[filterDuds, :],
1704
                allowedCharTimes[filterDuds, :],
1705
            )
1706

1707
            return sInds, intTime, slewTime, dV
×
1708

1709
        else:
1710
            empty = np.asarray([], dtype=int)
×
1711
            return empty, empty << u.d, empty << u.d, empty << self.m_per_s
×
1712

1713
    def chooseOcculterSlewTimes(self, sInds, slewTimes, dV, intTimes, charTimes):
1✔
1714
        """Selects the best slew time for each star
1715

1716
        This method searches through an array of permissible slew times for
1717
        each star and chooses the best slew time for the occulter based on
1718
        maximizing possible characterization time for that particular star (as
1719
        a default).
1720

1721
        Args:
1722
            sInds (~numpy.ndarray(int)):
1723
                Indices of available targets
1724
            slewTimes (astropy quantity array):
1725
                slew times to all stars (must be indexed by sInds)
1726
            dV (~astropy.units.Quantity(~numpy.ndarray(float))):
1727
                Delta-V used to transfer to new star line of sight in unis of m/s
1728
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1729
                Integration times for detection in units of day
1730
            charTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1731
                Time left over after integration which could be used for
1732
                characterization in units of day
1733

1734
        Returns:
1735
            tuple:
1736
                sInds (int):
1737
                    Indeces of next target star
1738
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1739
                    slew times to all stars (must be indexed by sInds)
1740
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1741
                    Integration times for detection in units of day
1742
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1743
                    Delta-V used to transfer to new star line of sight in unis of m/s
1744
        """
1745

1746
        # selection criteria for each star slew
1747
        good_j = np.argmax(
×
1748
            charTimes, axis=1
1749
        )  # maximum possible characterization time available
1750
        good_i = np.arange(0, len(sInds))
×
1751

1752
        dV = dV[good_i, good_j]
×
1753
        intTime = intTimes[good_i, good_j]
×
1754
        slewTime = slewTimes[good_i, good_j]
×
1755

1756
        return sInds, intTime, slewTime, dV
×
1757

1758
    def observation_detection(self, sInd, intTime, mode):
1✔
1759
        """Determines SNR and detection status for a given integration time
1760
        for detection. Also updates the lastDetected and starRevisit lists.
1761

1762
        Args:
1763
            sInd (int):
1764
                Integer index of the star of interest
1765
            intTime (~astropy.units.Quantity(~numpy.ndarray(float))):
1766
                Selected star integration time for detection in units of day.
1767
                Defaults to None.
1768
            mode (dict):
1769
                Selected observing mode for detection
1770

1771
        Returns:
1772
            tuple:
1773
                detected (numpy.ndarray(int)):
1774
                    Detection status for each planet orbiting the observed target star:
1775
                    1 is detection, 0 missed detection, -1 below IWA, and -2 beyond OWA
1776
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
1777
                    Surface brightness of local zodiacal light in units of 1/arcsec2
1778
                JEZ (astropy.units.Quantity(numpy.ndarray(float))):
1779
                    Intensity of exo-zodiacal light in units of photons/s/m2/arcsec2
1780
                systemParams (dict):
1781
                    Dictionary of time-dependant planet properties averaged over the
1782
                    duration of the integration
1783
                SNR (numpy.darray(float)):
1784
                    Detection signal-to-noise ratio of the observable planets
1785
                FA (bool):
1786
                    False alarm (false positive) boolean
1787

1788
        """
1789

1790
        PPop = self.PlanetPopulation
1✔
1791
        ZL = self.ZodiacalLight
1✔
1792
        PPro = self.PostProcessing
1✔
1793
        TL = self.TargetList
1✔
1794
        SU = self.SimulatedUniverse
1✔
1795
        Obs = self.Observatory
1✔
1796
        TK = self.TimeKeeping
1✔
1797

1798
        # Save Current Time before attempting time allocation
1799
        currentTimeNorm = TK.currentTimeNorm.copy()
1✔
1800
        currentTimeAbs = TK.currentTimeAbs.copy()
1✔
1801

1802
        # Allocate Time
1803
        extraTime = intTime * (mode["timeMultiplier"] - 1.0)  # calculates extraTime
1✔
1804
        success = TK.allocate_time(
1✔
1805
            intTime + extraTime + Obs.settlingTime + mode["syst"]["ohTime"], True
1806
        )
1807
        assert success, "Could not allocate observation detection time ({}).".format(
1✔
1808
            intTime + extraTime + Obs.settlingTime + mode["syst"]["ohTime"]
1809
        )
1810
        # calculates partial time to be added for every ntFlux
1811
        dt = intTime / float(self.ntFlux)
1✔
1812
        # find indices of planets around the target
1813
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
1814

1815
        # initialize outputs
1816
        detected = np.array([], dtype=int)
1✔
1817
        # write current system params by default
1818
        systemParams = SU.dump_system_params(sInd)
1✔
1819
        SNR = np.zeros(len(pInds))
1✔
1820

1821
        # if any planet, calculate SNR
1822
        if len(pInds) > 0:
1✔
1823
            # initialize arrays for SNR integration
1824
            fZs = np.zeros(self.ntFlux) << self.fZ_unit
1✔
1825
            systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
1826
            JEZs = np.zeros((self.ntFlux, len(pInds))) << self.JEZ_unit
1✔
1827
            Ss = np.zeros((self.ntFlux, len(pInds)))
1✔
1828
            Ns = np.zeros((self.ntFlux, len(pInds)))
1✔
1829
            # accounts for the time since the current time
1830
            timePlus = Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
1✔
1831
            # integrate the signal (planet flux) and noise
1832
            for i in range(self.ntFlux):
1✔
1833
                # allocate first half of dt
1834
                timePlus += dt / 2.0
1✔
1835
                # calculate current zodiacal light brightness
1836
                fZs[i] = ZL.fZ(
1✔
1837
                    Obs,
1838
                    TL,
1839
                    np.array([sInd], ndmin=1),
1840
                    (currentTimeAbs + timePlus).reshape(1),
1841
                    mode,
1842
                )[0]
1843
                # propagate the system to match up with current time
1844
                SU.propag_system(
1✔
1845
                    sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
1846
                )
1847
                # Calculate the exozodi intensity
1848
                JEZs[i] = SU.scale_JEZ(sInd, mode)
1✔
1849
                self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
1850
                # save planet parameters
1851
                systemParamss[i] = SU.dump_system_params(sInd)
1✔
1852
                # calculate signal and noise (electron count rates)
1853
                Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
1854
                    sInd, pInds, dt, mode, fZ=fZs[i], JEZ=JEZs[i]
1855
                )
1856
                # allocate second half of dt
1857
                timePlus += dt / 2.0
1✔
1858

1859
            # average output parameters
1860
            fZ = np.mean(fZs)
1✔
1861
            JEZ = np.mean(JEZs, axis=0)
1✔
1862
            systemParams = {
1✔
1863
                key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
1864
                / float(self.ntFlux)
1865
                for key in sorted(systemParamss[0])
1866
            }
1867
            # calculate SNR
1868
            S = Ss.sum(0)
1✔
1869
            N = Ns.sum(0)
1✔
1870
            SNR[N > 0] = S[N > 0] / N[N > 0]
1✔
1871

1872
        # if no planet, just save zodiacal brightness in the middle of the integration
1873
        else:
1874
            totTime = intTime * (mode["timeMultiplier"])
1✔
1875
            fZ = ZL.fZ(
1✔
1876
                Obs,
1877
                TL,
1878
                np.array([sInd], ndmin=1),
1879
                (currentTimeAbs + totTime / 2.0).reshape(1),
1880
                mode,
1881
            )[0]
1882
            # Use the default star value if no planets
1883
            JEZ = TL.JEZ0[mode["hex"]][sInd]
1✔
1884

1885
        # find out if a false positive (false alarm) or any false negative
1886
        # (missed detections) have occurred
1887
        FA, MD = PPro.det_occur(SNR, mode, TL, sInd, intTime)
1✔
1888

1889
        # populate detection status array
1890
        # 1:detected, 0:missed, -1:below IWA, -2:beyond OWA
1891
        if len(pInds) > 0:
1✔
1892
            detected = (~MD).astype(int)
1✔
1893
            WA = (
1✔
1894
                np.array(
1895
                    [
1896
                        systemParamss[x]["WA"].to_value(u.arcsec)
1897
                        for x in range(len(systemParamss))
1898
                    ]
1899
                )
1900
                << u.arcsec
1901
            )
1902
            detected[np.all(WA < mode["IWA"], 0)] = -1
1✔
1903
            detected[np.all(WA > mode["OWA"], 0)] = -2
1✔
1904

1905
        # if planets are detected, calculate the minimum apparent separation
1906
        smin = None
1✔
1907
        det = detected == 1  # If any of the planets around the star have been detected
1✔
1908
        if np.any(det):
1✔
1909
            smin = np.min(SU.s[pInds[det]])
1✔
1910
            log_det = "   - Detected planet inds %s (%s/%s)" % (
1✔
1911
                pInds[det],
1912
                len(pInds[det]),
1913
                len(pInds),
1914
            )
1915
            self.logger.info(log_det)
1✔
1916
            self.vprint(log_det)
1✔
1917

1918
        # populate the lastDetected array by storing det, JEZ, dMag, and WA
1919
        self.lastDetected[sInd, :] = [
1✔
1920
            det,
1921
            JEZ.flatten(),
1922
            systemParams["dMag"],
1923
            systemParams["WA"].to("arcsec").value,
1924
        ]
1925

1926
        # in case of a FA, generate a random delta mag (between PPro.FAdMag0 and
1927
        # TL.intCutoff_dMag) and working angle (between IWA and min(OWA, a_max))
1928
        if FA:
1✔
1929
            WA = (
×
1930
                np.random.uniform(
1931
                    mode["IWA"].to_value(u.arcsec),
1932
                    np.minimum(
1933
                        mode["OWA"], np.arctan(max(PPop.arange) / TL.dist[sInd])
1934
                    ).to_value(u.arcsec),
1935
                )
1936
                << u.arcsec
1937
            )
1938
            dMag = np.random.uniform(PPro.FAdMag0(WA), TL.intCutoff_dMag)
×
1939
            self.lastDetected[sInd, 0] = np.append(self.lastDetected[sInd, 0], True)
×
1940
            self.lastDetected[sInd, 1] = np.append(
×
1941
                self.lastDetected[sInd, 1], TL.JEZ0[mode["hex"]][sInd].flatten()
1942
            )
1943
            self.lastDetected[sInd, 2] = np.append(self.lastDetected[sInd, 2], dMag)
×
1944
            self.lastDetected[sInd, 3] = np.append(
×
1945
                self.lastDetected[sInd, 3], WA.to_value(u.arcsec)
1946
            )
1947
            sminFA = np.tan(WA) * TL.dist[sInd].to("AU")
×
1948
            smin = np.minimum(smin, sminFA) if smin is not None else sminFA
×
1949
            log_FA = "   - False Alarm (WA=%s, dMag=%s)" % (
×
1950
                np.round(WA, 3),
1951
                np.round(dMag, 1),
1952
            )
1953
            self.logger.info(log_FA)
×
1954
            self.vprint(log_FA)
×
1955

1956
        # Schedule Target Revisit
1957
        self.scheduleRevisit(sInd, smin, det, pInds)
1✔
1958

1959
        if self.make_debug_bird_plots:
1✔
1960
            from tools.obs_plot import obs_plot
×
1961

1962
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, detected)
×
1963

1964
        return detected.astype(int), fZ, JEZ, systemParams, SNR, FA
1✔
1965

1966
    def scheduleRevisit(self, sInd, smin, det, pInds):
1✔
1967
        """A Helper Method for scheduling revisits after observation detection
1968

1969
        Updates self.starRevisit attribute only
1970

1971
        Args:
1972
            sInd (int):
1973
                sInd of the star just detected
1974
            smin (~astropy.units.Quantity):
1975
                minimum separation of the planet to star of planet just detected
1976
            det (~np.ndarray(bool)):
1977
                Detection status of all planets in target system
1978
            pInds (~np.ndarray(int)):
1979
                Indices of planets in the target system
1980

1981
        Returns:
1982
            None
1983

1984
        """
1985
        TK = self.TimeKeeping
1✔
1986
        TL = self.TargetList
1✔
1987
        SU = self.SimulatedUniverse
1✔
1988

1989
        # in both cases (detection or false alarm), schedule a revisit
1990
        # based on minimum separation
1991
        Ms = TL.MsTrue[sInd]
1✔
1992
        if smin is not None:  # smin is None if no planet was detected
1✔
1993
            sp = smin
1✔
1994
            if np.any(det):
1✔
1995
                pInd_smin = pInds[det][np.argmin(SU.s[pInds[det]])]
1✔
1996
                Mp = SU.Mp[pInd_smin]
1✔
1997
            else:
1998
                Mp = SU.Mp.mean()
×
1999
            mu = const.G * (Mp + Ms)
1✔
2000
            T = 2.0 * np.pi * np.sqrt(sp**3.0 / mu)
1✔
2001
            t_rev = TK.currentTimeNorm.copy() + T / 2.0
1✔
2002
        # otherwise, revisit based on average of population semi-major axis and mass
2003
        else:
2004
            sp = SU.s.mean()
1✔
2005
            Mp = SU.Mp.mean()
1✔
2006
            mu = const.G * (Mp + Ms)
1✔
2007
            T = 2.0 * np.pi * np.sqrt(sp**3.0 / mu)
1✔
2008
            t_rev = TK.currentTimeNorm.copy() + 0.75 * T
1✔
2009

2010
        if self.revisit_wait is not None:
1✔
2011
            revisit_wait = self.revisit_wait
1✔
2012
            if np.ndim(revisit_wait) > 0:
1✔
2013
                revisit_wait = revisit_wait[sInd]
1✔
2014
            t_rev = TK.currentTimeNorm.copy() + revisit_wait
1✔
2015
        # finally, populate the revisit list (NOTE: sInd becomes a float)
2016
        revisit = np.array([sInd, float(t_rev.to_value(u.day))])
1✔
2017
        if self.starRevisit.size == 0:  # If starRevisit has nothing in it
1✔
2018
            self.starRevisit = np.array([revisit])  # initialize sterRevisit
1✔
2019
        else:
2020
            revInd = np.where(self.starRevisit[:, 0] == sInd)[
1✔
2021
                0
2022
            ]  # indices of the first column of the starRevisit list containing sInd
2023
            if revInd.size == 0:
1✔
2024
                self.starRevisit = np.vstack((self.starRevisit, revisit))
1✔
2025
            else:
UNCOV
2026
                self.starRevisit[revInd, 1] = revisit[1]
×
2027

2028
    def observation_characterization(self, sInd, mode):
1✔
2029
        """Finds if characterizations are possible and relevant information
2030

2031
        Args:
2032
            sInd (int):
2033
                Integer index of the star of interest
2034
            mode (dict):
2035
                Selected observing mode for characterization
2036

2037
        Returns:
2038
            tuple:
2039
                characterized (list(int)):
2040
                    Characterization status for each planet orbiting the observed
2041
                    target star including False Alarm if any, where 1 is full spectrum,
2042
                    -1 partial spectrum, and 0 not characterized
2043
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
2044
                    Surface brightness of local zodiacal light in units of 1/arcsec2
2045
                JEZ (astropy.units.Quantity(numpy.ndarray(float))):
2046
                    Intensity of exo-zodiacal light in units of photons/s/m2/arcsec
2047
                systemParams (dict):
2048
                    Dictionary of time-dependant planet properties averaged over the
2049
                    duration of the integration
2050
                SNR (numpy.ndarray(float)):
2051
                    Characterization signal-to-noise ratio of the observable planets.
2052
                    Defaults to None.
2053
                intTime (astropy.units.Quantity(numpy.ndarray(float))):
2054
                    Selected star characterization time in units of day.
2055
                    Defaults to None.
2056

2057
        """
2058

2059
        OS = self.OpticalSystem
1✔
2060
        ZL = self.ZodiacalLight
1✔
2061
        TL = self.TargetList
1✔
2062
        SU = self.SimulatedUniverse
1✔
2063
        Obs = self.Observatory
1✔
2064
        TK = self.TimeKeeping
1✔
2065

2066
        # selecting appropriate koMap
2067
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
2068

2069
        # find indices of planets around the target
2070
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2071

2072
        # get the detected status, and check if there was a FA
2073
        det = self.lastDetected[sInd, 0]
1✔
2074
        FA = len(det) == len(pInds) + 1
1✔
2075
        if FA:
1✔
2076
            pIndsDet = np.append(pInds, -1)[det]
×
2077
        else:
2078
            pIndsDet = pInds[det]
1✔
2079

2080
        # initialize outputs, and check if there's anything (planet or FA)
2081
        # to characterize
2082
        characterized = np.zeros(len(det), dtype=int)
1✔
2083
        fZ = 0.0 * self.inv_arcsec2
1✔
2084
        JEZ = 0.0 * self.JEZ_unit
1✔
2085
        # write current system params by default
2086
        systemParams = SU.dump_system_params(sInd)
1✔
2087
        SNR = np.zeros(len(det))
1✔
2088
        intTime = None
1✔
2089
        if len(det) == 0:  # nothing to characterize
1✔
2090
            return characterized, fZ, JEZ, systemParams, SNR, intTime
1✔
2091

2092
        # look for last detected planets that have not been fully characterized
2093
        if not (FA):  # only true planets, no FA
1✔
2094
            tochar = self.fullSpectra[pIndsDet] == 0
1✔
2095
        else:  # mix of planets and a FA
2096
            truePlans = pIndsDet[:-1]
×
2097
            tochar = np.append((self.fullSpectra[truePlans] == 0), True)
×
2098

2099
        # 1/ find spacecraft orbital START position including overhead time,
2100
        # and check keepout angle
2101
        if np.any(tochar):
1✔
2102
            # start times
2103
            startTime = (
1✔
2104
                TK.currentTimeAbs.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2105
            )
2106

2107
            # if we're beyond mission end, bail out
2108
            if startTime >= TK.missionFinishAbs:
1✔
2109
                return characterized, fZ, systemParams, SNR, intTime
×
2110

2111
            startTimeNorm = (
1✔
2112
                TK.currentTimeNorm.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2113
            )
2114
            # planets to characterize
2115
            koTimeInd = np.where(np.round(startTime.value) - self.koTimes.value == 0)[
1✔
2116
                0
2117
            ][
2118
                0
2119
            ]  # find indice where koTime is startTime[0]
2120
            # wherever koMap is 1, the target is observable
2121
            tochar[tochar] = koMap[sInd][koTimeInd]
1✔
2122

2123
        # 2/ if any planet to characterize, find the characterization times
2124
        # at the detected JEZ, dMag, and WA
2125
        if np.any(tochar):
1✔
2126
            fZ = ZL.fZ(Obs, TL, np.array([sInd], ndmin=1), startTime.reshape(1), mode)[
1✔
2127
                0
2128
            ]
2129
            JEZ = self.lastDetected[sInd, 1][det][tochar]
1✔
2130
            dMag = self.lastDetected[sInd, 2][det][tochar]
1✔
2131
            WA = self.lastDetected[sInd, 3][det][tochar] * u.arcsec
1✔
2132
            intTimes = np.zeros(len(tochar)) << u.day
1✔
2133
            intTimes[tochar] = OS.calc_intTime(TL, sInd, fZ, JEZ, dMag, WA, mode, TK=TK)
1✔
2134
            intTimes[~np.isfinite(intTimes)] = self.zero_d
1✔
2135
            # add a predetermined margin to the integration times
2136
            intTimes = intTimes * (1.0 + self.charMargin)
1✔
2137
            # apply time multiplier
2138
            totTimes = intTimes * (mode["timeMultiplier"])
1✔
2139

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

2144
            # end times
2145
            endTimes = startTime + totTimes
1✔
2146
            endTimesNorm = startTimeNorm + totTimes
1✔
2147
            # planets to characterize
2148
            tochar = (
1✔
2149
                (totTimes > 0)
2150
                & (totTimes <= OS.intCutoff)
2151
                & (endTimesNorm <= TK.OBendTimes[TK.OBnumber])
2152
            )
2153

2154
        # 3/ is target still observable at the end of any char time?
2155
        if np.any(tochar) and Obs.checkKeepoutEnd:
1✔
2156
            koTimeInds = np.zeros(len(endTimes.value[tochar]), dtype=int)
1✔
2157
            # find index in koMap where each endTime is closest to koTimes
2158
            for t, endTime in enumerate(endTimes.value[tochar]):
1✔
2159
                if endTime > self.koTimes.value[-1]:
1✔
2160
                    # case where endTime exceeds largest koTimes element
2161
                    endTimeInBounds = np.where(
×
2162
                        np.floor(endTime) - self.koTimes.value == 0
2163
                    )[0]
2164
                    koTimeInds[t] = (
×
2165
                        endTimeInBounds[0] if endTimeInBounds.size != 0 else -1
2166
                    )
2167
                else:
2168
                    koTimeInds[t] = np.where(
1✔
2169
                        np.round(endTime) - self.koTimes.value == 0
2170
                    )[0][
2171
                        0
2172
                    ]  # find indice where koTime is endTimes[0]
2173
            tochar[tochar] = [koMap[sInd][koT] if koT >= 0 else 0 for koT in koTimeInds]
1✔
2174

2175
        # 4/ if yes, allocate the overhead time, and perform the characterization
2176
        # for the maximum char time
2177
        if np.any(tochar):
1✔
2178
            # Save Current Time before attempting time allocation
2179
            currentTimeNorm = TK.currentTimeNorm.copy()
1✔
2180
            currentTimeAbs = TK.currentTimeAbs.copy()
1✔
2181

2182
            # Allocate Time
2183
            intTime = np.max(intTimes[tochar])
1✔
2184
            extraTime = intTime * (mode["timeMultiplier"] - 1.0)
1✔
2185
            success = TK.allocate_time(
1✔
2186
                intTime + extraTime + mode["syst"]["ohTime"] + Obs.settlingTime, True
2187
            )  # allocates time
2188
            if not (success):  # Time was not successfully allocated
1✔
UNCOV
2189
                char_intTime = None
×
UNCOV
2190
                lenChar = len(pInds) + 1 if FA else len(pInds)
×
UNCOV
2191
                characterized = np.zeros(lenChar, dtype=float)
×
UNCOV
2192
                char_JEZ = np.zeros(lenChar, dtype=float) << self.JEZ_unit
×
UNCOV
2193
                char_SNR = np.zeros(lenChar, dtype=float)
×
UNCOV
2194
                char_fZ = 0.0 * self.inv_arcsec2
×
UNCOV
2195
                char_systemParams = SU.dump_system_params(sInd)
×
UNCOV
2196
                return (
×
2197
                    characterized,
2198
                    char_fZ,
2199
                    char_JEZ,
2200
                    char_systemParams,
2201
                    char_SNR,
2202
                    char_intTime,
2203
                )
2204

2205
            pIndsChar = pIndsDet[tochar]
1✔
2206
            log_char = "   - Charact. planet inds %s (%s/%s detected)" % (
1✔
2207
                pIndsChar,
2208
                len(pIndsChar),
2209
                len(pIndsDet),
2210
            )
2211
            self.logger.info(log_char)
1✔
2212
            self.vprint(log_char)
1✔
2213

2214
            # SNR CALCULATION:
2215
            # first, calculate SNR for observable planets (without false alarm)
2216
            planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
1✔
2217
            SNRplans = np.zeros(len(planinds))
1✔
2218
            if len(planinds) > 0:
1✔
2219
                # initialize arrays for SNR integration
2220
                fZs = np.zeros(self.ntFlux) << self.inv_arcsec2
1✔
2221
                systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
2222
                JEZs = np.zeros((self.ntFlux, len(planinds))) << self.JEZ_unit
1✔
2223
                Ss = np.zeros((self.ntFlux, len(planinds)))
1✔
2224
                Ns = np.zeros((self.ntFlux, len(planinds)))
1✔
2225
                # integrate the signal (planet flux) and noise
2226
                dt = intTime / float(self.ntFlux)
1✔
2227
                timePlus = (
1✔
2228
                    Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2229
                )  # accounts for the time since the current time
2230
                for i in range(self.ntFlux):
1✔
2231
                    # allocate first half of dt
2232
                    timePlus += dt / 2.0
1✔
2233
                    # calculate current zodiacal light brightness
2234
                    fZs[i] = ZL.fZ(
1✔
2235
                        Obs,
2236
                        TL,
2237
                        np.array([sInd], ndmin=1),
2238
                        (currentTimeAbs + timePlus).reshape(1),
2239
                        mode,
2240
                    )[0]
2241
                    # propagate the system to match up with current time
2242
                    SU.propag_system(
1✔
2243
                        sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2244
                    )
2245
                    self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
2246
                    # Calculate the exozodi intensity
2247
                    JEZs[i] = SU.scale_JEZ(sInd, mode, pInds=planinds)
1✔
2248
                    # save planet parameters
2249
                    systemParamss[i] = SU.dump_system_params(sInd)
1✔
2250
                    # calculate signal and noise (electron count rates)
2251
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
2252
                        sInd, planinds, dt, mode, fZ=fZs[i], JEZ=JEZs[i]
2253
                    )
2254
                    # allocate second half of dt
2255
                    timePlus += dt / 2.0
1✔
2256

2257
                # average output parameters
2258
                fZ = np.mean(fZs)
1✔
2259
                JEZ = np.mean(JEZs, axis=0)
1✔
2260
                systemParams = {
1✔
2261
                    key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2262
                    / float(self.ntFlux)
2263
                    for key in sorted(systemParamss[0])
2264
                }
2265
                # calculate planets SNR
2266
                S = Ss.sum(0)
1✔
2267
                N = Ns.sum(0)
1✔
2268
                SNRplans[N > 0] = S[N > 0] / N[N > 0]
1✔
2269
                # allocate extra time for timeMultiplier
2270

2271
            # if only a FA, just save zodiacal brightness in the middle of the
2272
            # integration
2273
            else:
2274
                totTime = intTime * (mode["timeMultiplier"])
×
2275
                fZ = ZL.fZ(
×
2276
                    Obs,
2277
                    TL,
2278
                    np.array([sInd], ndmin=1),
2279
                    (currentTimeAbs + totTime / 2.0).reshape(1),
2280
                    mode,
2281
                )[0]
2282
                # Use the default star value if no planets
2283
                JEZ = TL.JEZ0[mode["hex"]][sInd]
×
2284

2285
            # calculate the false alarm SNR (if any)
2286
            SNRfa = []
1✔
2287
            if pIndsChar[-1] == -1:
1✔
2288
                JEZ = self.lastDetected[sInd, 1][-1]
×
2289
                dMag = self.lastDetected[sInd, 2][-1]
×
2290
                WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
2291
                C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, JEZ, dMag, WA, mode, TK=TK)
×
2292
                S = (C_p * intTime).decompose().value
×
2293
                N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2294
                SNRfa = S / N if N > 0.0 else 0.0
×
2295

2296
            # save all SNRs (planets and FA) to one array
2297
            SNRinds = np.where(det)[0][tochar]
1✔
2298
            SNR[SNRinds] = np.append(SNRplans, SNRfa)
1✔
2299

2300
            # now, store characterization status: 1 for full spectrum,
2301
            # -1 for partial spectrum, 0 for not characterized
2302
            char = SNR >= mode["SNR"]
1✔
2303
            # initialize with full spectra
2304
            characterized = char.astype(int)
1✔
2305
            WAchar = self.lastDetected[sInd, 3][char] * u.arcsec
1✔
2306
            # find the current WAs of characterized planets
2307
            WAs = systemParams["WA"]
1✔
2308
            if FA:
1✔
2309
                WAs = np.append(WAs, self.lastDetected[sInd, 3][-1] * u.arcsec)
×
2310
            # check for partial spectra (for coronagraphs only)
2311
            if not (mode["syst"]["occulter"]):
1✔
2312
                IWA_max = mode["IWA"] * (1.0 + mode["BW"] / 2.0)
1✔
2313
                OWA_min = mode["OWA"] * (1.0 - mode["BW"] / 2.0)
1✔
2314
                char[char] = (WAchar < IWA_max) | (WAchar > OWA_min)
1✔
2315
                characterized[char] = -1
1✔
2316
            # encode results in spectra lists (only for planets, not FA)
2317
            charplans = characterized[:-1] if FA else characterized
1✔
2318
            self.fullSpectra[pInds[charplans == 1]] += 1
1✔
2319
            self.partialSpectra[pInds[charplans == -1]] += 1
1✔
2320

2321
        if self.make_debug_bird_plots:
1✔
2322
            from tools.obs_plot import obs_plot
×
2323

2324
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2325

2326
        return characterized.astype(int), fZ, JEZ, systemParams, SNR, intTime
1✔
2327

2328
    def calc_signal_noise(
1✔
2329
        self, sInd, pInds, t_int, mode, fZ=None, JEZ=None, dMag=None, WA=None
2330
    ):
2331
        """Calculates the signal and noise fluxes for a given time interval. Called
2332
        by observation_detection and observation_characterization methods in the
2333
        SurveySimulation module.
2334

2335
        Args:
2336
            sInd (int):
2337
                Integer index of the star of interest
2338
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2339
                Integration time interval in units of day
2340
            pInds (int):
2341
                Integer indices of the planets of interest
2342
            mode (dict):
2343
                Selected observing mode (from OpticalSystem)
2344
            fZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2345
                Surface brightness of local zodiacal light in units of 1/arcsec2
2346
            JEZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2347
                Intensity of exo-zodiacal light in units of ph/s/m2/arcsec2
2348
            dMag (~numpy.ndarray(float)):
2349
                Differences in magnitude between planets and their host star
2350
            WA (~astropy.units.Quantity(~numpy.ndarray(float))):
2351
                Working angles of the planets of interest in units of arcsec
2352

2353
        Returns:
2354
            tuple:
2355
                Signal (float):
2356
                    Counts of signal
2357
                Noise (float):
2358
                    Counts of background noise variance
2359

2360
        """
2361

2362
        OS = self.OpticalSystem
1✔
2363
        ZL = self.ZodiacalLight
1✔
2364
        TL = self.TargetList
1✔
2365
        SU = self.SimulatedUniverse
1✔
2366
        Obs = self.Observatory
1✔
2367
        TK = self.TimeKeeping
1✔
2368

2369
        # calculate optional parameters if not provided
2370
        fZ = (
1✔
2371
            fZ
2372
            if (fZ is not None)
2373
            else ZL.fZ(
2374
                Obs,
2375
                TL,
2376
                np.array([sInd], ndmin=1),
2377
                TK.currentTimeAbs.copy().reshape(1),
2378
                mode,
2379
            )[0]
2380
        )
2381
        JEZ = (
1✔
2382
            u.Quantity(JEZ, ndmin=1)
2383
            if (JEZ is not None)
2384
            else SU.scale_JEZ(sInd, mode, pInds=pInds)
2385
        )
2386

2387
        # if lucky_planets, use lucky planet params for dMag and WA
2388
        if SU.lucky_planets and mode in list(
1✔
2389
            filter(lambda mode: "spec" in mode["inst"]["name"], OS.observingModes)
2390
        ):
2391
            phi = (1 / np.pi) * np.ones(len(SU.d))
×
2392
            dMag = deltaMag(SU.p, SU.Rp, SU.d, phi)[pInds]  # delta magnitude
×
2393
            # working angle
2394
            WA = (
×
2395
                np.arctan(
2396
                    SU.a.to_value(u.AU)
2397
                    * self.AU2pc
2398
                    / TL.dist[SU.plan2star].to_value(u.pc)
2399
                )
2400
                * self.rad2arcsec
2401
                << u.arcsec
2402
            )[pInds]
2403
        else:
2404
            dMag = dMag if (dMag is not None) else SU.dMag[pInds]
1✔
2405
            WA = WA if (WA is not None) else SU.WA[pInds]
1✔
2406

2407
        # initialize Signal and Noise arrays
2408
        Signal = np.zeros(len(pInds))
1✔
2409
        Noise = np.zeros(len(pInds))
1✔
2410

2411
        # find observable planets wrt IWA-OWA range
2412
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2413

2414
        if np.any(obs):
1✔
2415
            # find electron counts
2416
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(
1✔
2417
                TL, sInd, fZ, JEZ[obs], dMag[obs], WA[obs], mode, TK=TK
2418
            )
2419
            # calculate signal and noise levels (based on Nemati14 formula)
2420
            _t_int_s = t_int.to_value(u.d) * self.day2sec
1✔
2421

2422
            Signal[obs] = C_p.to_value(self.inv_s) * _t_int_s
1✔
2423
            Noise[obs] = np.sqrt(
1✔
2424
                (
2425
                    C_b.to_value(self.inv_s) * _t_int_s
2426
                    + (C_sp.to_value(self.inv_s) * _t_int_s) ** 2
2427
                )
2428
            )
2429

2430
        return Signal, Noise
1✔
2431

2432
    def update_occulter_mass(self, DRM, sInd, t_int, skMode):
1✔
2433
        """Updates the occulter wet mass in the Observatory module, and stores all
2434
        the occulter related values in the DRM array.
2435

2436
        Args:
2437
            DRM (dict):
2438
                Design Reference Mission, contains the results of one complete
2439
                observation (detection and characterization)
2440
            sInd (int):
2441
                Integer index of the star of interest
2442
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2443
                Selected star integration time (for detection or characterization)
2444
                in units of day
2445
            skMode (str):
2446
                Station keeping observing mode type ('det' or 'char')
2447

2448
        Returns:
2449
            dict:
2450
                Design Reference Mission dictionary, contains the results of one
2451
                complete observation
2452

2453
        """
2454

2455
        TL = self.TargetList
×
2456
        Obs = self.Observatory
×
2457
        TK = self.TimeKeeping
×
2458

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

2461
        # decrement mass for station-keeping
2462
        dF_lateral, dF_axial, intMdot, mass_used, deltaV = Obs.mass_dec_sk(
×
2463
            TL, sInd, TK.currentTimeAbs.copy(), t_int
2464
        )
2465

2466
        DRM[skMode + "_dV"] = deltaV.to("m/s")
×
2467
        DRM[skMode + "_mass_used"] = mass_used.to("kg")
×
2468
        DRM[skMode + "_dF_lateral"] = dF_lateral.to("N")
×
2469
        DRM[skMode + "_dF_axial"] = dF_axial.to("N")
×
2470
        # update spacecraft mass
2471
        Obs.scMass = Obs.scMass - mass_used
×
2472
        DRM["scMass"] = Obs.scMass.to("kg")
×
2473
        if Obs.twotanks:
×
2474
            Obs.skMass = Obs.skMass - mass_used
×
2475
            DRM["skMass"] = Obs.skMass.to("kg")
×
2476

2477
        return DRM
×
2478

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

2482
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2483
        objects with their own outspecs.
2484

2485
        Args:
2486
            genNewPlanets (bool):
2487
                Generate all new planets based on the original input specification.
2488
                If False, then the original planets will remain. Setting to True forces
2489
                ``rewindPlanets`` to be True as well. Defaults True.
2490
            rewindPlanets (bool):
2491
                Reset the current set of planets to their original orbital phases.
2492
                If both genNewPlanets and rewindPlanet  are False, the original planets
2493
                will be retained and will not be rewound to their initial starting
2494
                locations (i.e., all systems will remain at the times they were at the
2495
                end of the last run, thereby effectively randomizing planet phases.
2496
                Defaults True.
2497
            seed (int, optional):
2498
                Random seed to use for all random number generation. If None (default)
2499
                a new random seed will be generated when re-initializing the
2500
                SurveySimulation.
2501

2502
        """
2503

2504
        SU = self.SimulatedUniverse
1✔
2505
        TK = self.TimeKeeping
1✔
2506
        TL = self.TargetList
1✔
2507

2508
        # re-initialize SurveySimulation arrays
2509
        specs = self._outspec
1✔
2510
        specs["modules"] = self.modules
1✔
2511

2512
        if seed is None:  # pop the seed so a new one is generated
1✔
2513
            if "seed" in specs:
1✔
2514
                specs.pop("seed")
1✔
2515
        else:  # if seed is provided, replace seed with provided seed
2516
            specs["seed"] = seed
×
2517

2518
        # reset mission time, re-init surveysim and observatory
2519
        TK.__init__(**TK._outspec)
1✔
2520
        self.__init__(**specs)
1✔
2521
        self.Observatory.__init__(**self.Observatory._outspec)
1✔
2522

2523
        # generate new planets if requested (default)
2524
        if genNewPlanets:
1✔
2525
            TL.stellar_mass()
1✔
2526
            TL.I = TL.gen_inclinations(TL.PlanetPopulation.Irange)  # noqa: E741
1✔
2527
            SU.gen_physical_properties(**SU._outspec)
1✔
2528
            rewindPlanets = True
1✔
2529
        # re-initialize systems if requested (default)
2530
        if rewindPlanets:
1✔
2531
            SU.init_systems()
1✔
2532

2533
        # reset helper arrays
2534
        self.initializeStorageArrays()
1✔
2535

2536
        self.vprint("Simulation reset.")
1✔
2537

2538
    def genOutSpec(
1✔
2539
        self,
2540
        tofile: Optional[str] = None,
2541
        starting_outspec: Optional[Dict[str, Any]] = None,
2542
        modnames: bool = False,
2543
    ) -> Dict[str, Any]:
2544
        """Join all _outspec dicts from all modules into one output dict
2545
        and optionally write out to JSON file on disk.
2546

2547
        Args:
2548
            tofile (str, optional):
2549
                Name of the file containing all output specifications (outspecs).
2550
                Defaults to None.
2551
            starting_outspec (dict, optional):
2552
                Initial outspec (from MissionSim). Defaults to None.
2553
            modnames (bool):
2554
                If True, populate outspec dictionary with the module it originated from,
2555
                instead of the actual value of the keyword. Defaults False.
2556

2557
        Returns:
2558
            dict:
2559
                Dictionary containing the full :ref:`sec:inputspec`, including all
2560
                filled-in default values. Combination of all individual module _outspec
2561
                attributes.
2562

2563
        """
2564

2565
        # start with our own outspec
2566
        if modnames:
1✔
2567
            out = copy.copy(self._outspec)
1✔
2568
            for k in out:
1✔
2569
                out[k] = self.__class__.__name__
1✔
2570
        else:
2571
            out = copy.deepcopy(self._outspec)
1✔
2572

2573
        # Add any provided other outspec
2574
        if starting_outspec is not None:
1✔
2575
            out.update(starting_outspec)
1✔
2576

2577
        # add in all modules _outspec's
2578
        for module in self.modules.values():
1✔
2579
            if modnames:
1✔
2580
                tmp = copy.copy(module._outspec)
1✔
2581
                for k in tmp:
1✔
2582
                    tmp[k] = module.__class__.__name__
1✔
2583
            else:
2584
                tmp = module._outspec
1✔
2585
            out.update(tmp)
1✔
2586

2587
        # add in the specific module names used
2588
        out["modules"] = {}
1✔
2589
        for mod_name, module in self.modules.items():
1✔
2590
            # find the module file
2591
            mod_name_full = module.__module__
1✔
2592
            if mod_name_full.startswith("EXOSIMS"):
1✔
2593
                # take just its short name if it is in EXOSIMS
2594
                mod_name_short = mod_name_full.split(".")[-1]
1✔
2595
            else:
2596
                # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2597
                mod_name_short = re.sub(
×
2598
                    r"\.pyc$", ".py", inspect.getfile(module.__class__)
2599
                )
2600
            out["modules"][mod_name] = mod_name_short
1✔
2601
        # add catalog name
2602
        module = self.TargetList.StarCatalog
1✔
2603
        mod_name_full = module.__module__
1✔
2604
        if mod_name_full.startswith("EXOSIMS"):
1✔
2605
            # take just its short name if it is in EXOSIMS
2606
            mod_name_short = mod_name_full.split(".")[-1]
1✔
2607
        else:
2608
            # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2609
            mod_name_short = re.sub(r"\.pyc$", ".py", inspect.getfile(module.__class__))
×
2610
        out["modules"]["StarCatalog"] = mod_name_short
1✔
2611

2612
        # if we don't know about the SurveyEnsemble, just write a blank to the output
2613
        if "SurveyEnsemble" not in out["modules"]:
1✔
2614
            out["modules"]["SurveyEnsemble"] = " "
1✔
2615

2616
        # get version and Git information
2617
        out["Version"] = get_version()
1✔
2618

2619
        # dump to file
2620
        if tofile is not None:
1✔
2621
            with open(tofile, "w") as outfile:
1✔
2622
                json.dump(
1✔
2623
                    out,
2624
                    outfile,
2625
                    sort_keys=True,
2626
                    indent=4,
2627
                    ensure_ascii=False,
2628
                    separators=(",", ": "),
2629
                    default=array_encoder,
2630
                )
2631

2632
        return out
1✔
2633

2634
    def generateHashfName(self, specs):
1✔
2635
        """Generate cached file Hashname
2636

2637
        Requires a .XXX appended to end of hashname for each individual use case
2638

2639
        Args:
2640
            specs (dict):
2641
                :ref:`sec:inputspec`
2642

2643
        Returns:
2644
            str:
2645
                Unique indentifier string for cahce products from this set of modules
2646
                and inputs
2647
        """
2648
        # Allows cachefname to be predefined
2649
        if "cachefname" in specs:
1✔
2650
            return specs["cachefname"]
×
2651

2652
        cachefname = ""  # declares cachefname
1✔
2653
        mods = ["Completeness", "TargetList", "OpticalSystem"]  # modules to look at
1✔
2654
        tmp = (
1✔
2655
            self.Completeness.PlanetPopulation.__class__.__name__
2656
            + self.Completeness.PlanetPhysicalModel.__class__.__name__
2657
            + self.PlanetPopulation.__class__.__name__
2658
            + self.SimulatedUniverse.__class__.__name__
2659
            + self.PlanetPhysicalModel.__class__.__name__
2660
        )
2661

2662
        if "selectionMetric" in specs:
1✔
2663
            tmp += specs["selectionMetric"]
×
2664
        if "Izod" in specs:
1✔
2665
            tmp += specs["Izod"]
×
2666
        if "maxiter" in specs:
1✔
2667
            tmp += str(specs["maxiter"])
×
2668
        if "ftol" in specs:
1✔
2669
            tmp += str(specs["ftol"])
×
2670
        if "missionLife" in specs:
1✔
2671
            tmp += str(specs["missionLife"])
1✔
2672
        if "missionPortion" in specs:
1✔
2673
            tmp += str(specs["missionPortion"])
1✔
2674
        if "smaknee" in specs:
1✔
2675
            tmp += str(specs["smaknee"])
×
2676
        if "koAngleMax" in specs:
1✔
2677
            tmp += str(specs["koAngleMax"])
×
2678
        tmp += str(np.sum(self.Completeness.PlanetPopulation.arange.value))
1✔
2679
        tmp += str(np.sum(self.Completeness.PlanetPopulation.Rprange.value))
1✔
2680
        tmp += str(np.sum(self.Completeness.PlanetPopulation.erange))
1✔
2681
        tmp += str(
1✔
2682
            self.Completeness.PlanetPopulation.PlanetPhysicalModel.whichPlanetPhaseFunction  # noqa: E501
2683
        )
2684
        tmp += str(np.sum(self.PlanetPopulation.arange.value))
1✔
2685
        tmp += str(np.sum(self.PlanetPopulation.Rprange.value))
1✔
2686
        tmp += str(np.sum(self.PlanetPopulation.erange))
1✔
2687
        tmp += str(self.PlanetPopulation.PlanetPhysicalModel.whichPlanetPhaseFunction)
1✔
2688

2689
        for mod in mods:
1✔
2690
            cachefname += self.modules[mod].__module__.split(".")[
1✔
2691
                -1
2692
            ]  # add module name to end of cachefname
2693
        cachefname += hashlib.md5(
1✔
2694
            (str(self.TargetList.Name) + tmp).encode("utf-8")
2695
        ).hexdigest()  # turn cachefname into hashlib
2696
        cachefname = os.path.join(
1✔
2697
            self.cachedir, cachefname + os.extsep
2698
        )  # join into filepath and fname
2699
        # Needs file terminator (.starkt0, .t0, etc) appended done by each individual
2700
        # use case.
2701
        return cachefname
1✔
2702

2703
    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
1✔
2704
        """Helper method for Overloading Revisit Filtering
2705

2706
        Args:
2707
            sInds (~numpy.ndarray(int)):
2708
                Indices of stars still in observation list
2709
            tmpCurrentTimeNorm (astropy.units.Quantity):
2710
                Normalized simulation time
2711

2712
        Returns:
2713
            ~numpy.ndarray(int):
2714
                indices of stars still in observation list
2715
        """
2716

2717
        tovisit = np.zeros(self.TargetList.nStars, dtype=bool)
1✔
2718
        if len(sInds) > 0:
1✔
2719
            # Check that no star has exceeded the number of revisits and the indicies
2720
            # of all considered stars have minimum number of observations
2721
            # This should prevent revisits so long as all stars have not
2722
            # been observed
2723
            tovisit[sInds] = (self.starVisits[sInds] == min(self.starVisits[sInds])) & (
1✔
2724
                self.starVisits[sInds] < self.nVisitsMax
2725
            )
2726
            if self.starRevisit.size != 0:
1✔
2727
                ind_rev = self.revisit_inds(sInds, tmpCurrentTimeNorm)
1✔
2728
                tovisit[ind_rev] = self.starVisits[ind_rev] < self.nVisitsMax
1✔
2729
            sInds = np.where(tovisit)[0]
1✔
2730
        return sInds
1✔
2731

2732
    def is_earthlike(self, plan_inds, sInd):
1✔
2733
        """Is the planet earthlike? Returns boolean array that's True for Earth-like
2734
        planets.
2735

2736

2737
        Args:
2738
            plan_inds(~numpy.ndarray(int)):
2739
                Planet indices
2740
            sInd (int):
2741
                Star index
2742

2743
        Returns:
2744
            ~numpy.ndarray(bool):
2745
                Array of same dimension as plan_inds input that's True for Earthlike
2746
                planets and false otherwise.
2747
        """
2748
        TL = self.TargetList
1✔
2749
        SU = self.SimulatedUniverse
1✔
2750
        PPop = self.PlanetPopulation
1✔
2751

2752
        # extract planet and star properties
2753
        Rp_plan = SU.Rp[plan_inds].value
1✔
2754
        L_star = TL.L[sInd]
1✔
2755
        if PPop.scaleOrbits:
1✔
2756
            a_plan = (SU.a[plan_inds] / np.sqrt(L_star)).value
×
2757
        else:
2758
            a_plan = (SU.a[plan_inds]).value
1✔
2759
        # Definition: planet radius (in earth radii) and solar-equivalent luminosity
2760
        # must be between the given bounds.
2761
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
1✔
2762
        # We use the numpy versions so that plan_ind can be a numpy vector.
2763
        return np.logical_and(
1✔
2764
            np.logical_and(Rp_plan >= Rp_plan_lo, Rp_plan <= 1.4),
2765
            np.logical_and(a_plan >= 0.95, a_plan <= 1.67),
2766
        )
2767

2768
    def find_known_plans(self):
1✔
2769
        """
2770
        Find and return list of known RV stars and list of stars with earthlike planets
2771
        based on info from David Plavchan dated 12/24/2018
2772
        """
2773
        TL = self.TargetList
×
2774
        SU = self.SimulatedUniverse
×
2775
        PPop = self.PlanetPopulation
×
2776
        L_star = TL.L[SU.plan2star]
×
2777

2778
        c = 28.4 * u.m / u.s
×
2779
        Mj = 317.8 * u.earthMass
×
2780
        Mpj = SU.Mp / Mj  # planet masses in jupiter mass units
×
2781
        Ms = TL.MsTrue[SU.plan2star]
×
2782
        Teff = TL.Teff[SU.plan2star]
×
2783
        mu = const.G * (SU.Mp + Ms)
×
2784
        T = (2.0 * np.pi * np.sqrt(SU.a**3 / mu)).to(u.yr)
×
2785
        e = SU.e
×
2786

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

2790
        K = (
×
2791
            (c / np.sqrt(1 - e[t_filt]))
2792
            * Mpj[t_filt]
2793
            * np.sin(SU.I[t_filt])
2794
            * Ms[t_filt] ** (-2 / 3)
2795
            * T[t_filt] ** (-1 / 3)
2796
        )
2797

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

2803
        if PPop.scaleOrbits:
×
2804
            a_plan = (SU.a / np.sqrt(L_star)).value
×
2805
        else:
2806
            a_plan = SU.a.value
×
2807

2808
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
×
2809

2810
        # pinds in habitable zone
2811
        a_filt = k_filt[np.where((a_plan[k_filt] > 0.95) & (a_plan[k_filt] < 1.67))[0]]
×
2812
        # rocky planets
2813
        r_filt = a_filt[
×
2814
            np.where(
2815
                (SU.Rp.value[a_filt] >= Rp_plan_lo[a_filt])
2816
                & (SU.Rp.value[a_filt] < 1.4)
2817
            )[0]
2818
        ]
2819
        self.known_earths = np.union1d(self.known_earths, r_filt).astype(int)
×
2820

2821
        # these are known_rv stars with earths around them
2822
        known_stars = np.unique(SU.plan2star[k_filt])  # these are known_rv stars
×
2823
        known_rocky = np.unique(SU.plan2star[r_filt])
×
2824

2825
        # if include_known_RV, then filter out all other sInds
2826
        if self.include_known_RV is not None:
×
2827
            HIP_sInds = np.where(np.isin(TL.Name, self.include_known_RV))[0]
×
2828
            known_stars = np.intersect1d(HIP_sInds, known_stars)
×
2829
            known_rocky = np.intersect1d(HIP_sInds, known_rocky)
×
2830
        return known_stars.astype(int), known_rocky.astype(int)
×
2831

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

2835
        The observation window (which includes settling and overhead times)
2836
        is a superset of the integration window (in which photons are collected).
2837

2838
        The observation window begins at startTime. The integration window
2839
        begins at startTime + Obs.settlingTime + mode["ohTime"],
2840
        and the integration itself has a duration of intTime.
2841

2842
        Args:
2843
            sInd (int):
2844
                Integer index of the star of interest
2845
            pIndsChar (int numpy.ndarray):
2846
                Observable planets to characterize. Place (-1) at end to put
2847
                False Alarm parameters at end of returned tuples.
2848
            startTime (astropy.units.Quantity):
2849
                Beginning of observation window in units of day.
2850
            intTime (astropy.units.Quantity):
2851
                Selected star characterization integration time in units of day.
2852
            mode (dict):
2853
                Observing mode for the characterization
2854

2855
        Returns:
2856
            tuple:
2857
                SNR (float numpy.ndarray):
2858
                    Characterization signal-to-noise ratio of the observable planets.
2859
                    Defaults to None. [TBD]
2860
                fZ (astropy.units.Quantity):
2861
                    Surface brightness of local zodiacal light in units of 1/arcsec2
2862
                systemParams (dict):
2863
                    Dictionary of time-dependent planet properties averaged over the
2864
                    duration of the integration.
2865

2866
        """
2867

2868
        OS = self.OpticalSystem
1✔
2869
        ZL = self.ZodiacalLight
1✔
2870
        TL = self.TargetList
1✔
2871
        SU = self.SimulatedUniverse
1✔
2872
        Obs = self.Observatory
1✔
2873
        TK = self.TimeKeeping
1✔
2874

2875
        # time at start of integration window
2876
        currentTimeNorm = startTime
1✔
2877
        currentTimeAbs = TK.missionStart + startTime
1✔
2878

2879
        # first, calculate SNR for observable planets (without false alarm)
2880
        planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
1✔
2881
        SNRplans = np.zeros(len(planinds))
1✔
2882
        if len(planinds) > 0:
1✔
2883
            # initialize arrays for SNR integration
2884
            fZs = np.zeros(self.ntFlux) << self.fZ_unit
1✔
2885
            systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
2886
            Ss = np.zeros((self.ntFlux, len(planinds)))
1✔
2887
            Ns = np.zeros((self.ntFlux, len(planinds)))
1✔
2888
            # integrate the signal (planet flux) and noise
2889
            dt = intTime / float(self.ntFlux)
1✔
2890
            timePlus = (
1✔
2891
                Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2892
            )  # accounts for the time since the current time
2893
            for i in range(self.ntFlux):
1✔
2894
                # calculate signal and noise (electron count rates)
2895
                if SU.lucky_planets:
1✔
2896
                    fZs[i] = ZL.fZ(
×
2897
                        Obs,
2898
                        TL,
2899
                        np.array([sInd], ndmin=1),
2900
                        currentTimeAbs.reshape(1),
2901
                        mode,
2902
                    )[0]
2903
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
×
2904
                        sInd, planinds, dt, mode, fZ=fZs[i]
2905
                    )
2906
                # allocate first half of dt
2907
                timePlus += dt / 2.0
1✔
2908
                # calculate current zodiacal light brightness
2909
                fZs[i] = ZL.fZ(
1✔
2910
                    Obs,
2911
                    TL,
2912
                    np.array([sInd], ndmin=1),
2913
                    (currentTimeAbs + timePlus).reshape(1),
2914
                    mode,
2915
                )[0]
2916
                # propagate the system to match up with current time
2917
                SU.propag_system(
1✔
2918
                    sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2919
                )
2920
                self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
2921
                # time-varying planet params (WA, dMag, phi, fEZ, d)
2922
                systemParamss[i] = SU.dump_system_params(sInd)
1✔
2923
                # calculate signal and noise (electron count rates)
2924
                if not SU.lucky_planets:
1✔
2925
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
2926
                        sInd, planinds, dt, mode, fZ=fZs[i]
2927
                    )
2928
                # allocate second half of dt
2929
                timePlus += dt / 2.0
1✔
2930

2931
            # average output parameters
2932
            fZ = np.mean(fZs)
1✔
2933
            systemParams = {
1✔
2934
                key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2935
                / float(self.ntFlux)
2936
                for key in sorted(systemParamss[0])
2937
            }
2938
            # calculate planets SNR
2939
            S = Ss.sum(0)
1✔
2940
            N = Ns.sum(0)
1✔
2941
            SNRplans[N > 0] = S[N > 0] / N[N > 0]
1✔
2942
            # allocate extra time for timeMultiplier
2943

2944
        # if only a FA, just save zodiacal brightness
2945
        # in the middle of the integration
2946
        else:
2947
            totTime = intTime * (mode["timeMultiplier"])
×
2948
            fZ = ZL.fZ(
×
2949
                Obs,
2950
                TL,
2951
                np.array([sInd], ndmin=1),
2952
                (currentTimeAbs.copy() + totTime / 2.0).reshape(1),
2953
                mode,
2954
            )[0]
2955

2956
        # calculate the false alarm SNR (if any)
2957
        SNRfa = []
1✔
2958
        if pIndsChar[-1] == -1:
1✔
2959
            # Note: these attributes may not exist for all schedulers
2960
            fEZ = self.lastDetected[sInd, 1][-1] << self.fZ_unit
×
2961
            dMag = self.lastDetected[sInd, 2][-1]
×
2962
            WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
2963
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, fEZ, dMag, WA, mode)
×
2964
            S = (C_p * intTime).decompose().value
×
2965
            N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2966
            SNRfa = S / N if N > 0.0 else 0.0
×
2967

2968
        # SNR vector is of length (#planets), plus 1 if FA
2969
        # This routine leaves SNR = 0 if unknown or not found
2970
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2971
        FA_present = pIndsChar[-1] == -1
1✔
2972
        # there will be one SNR for each entry of pInds_FA
2973
        pInds_FA = np.append(pInds, [-1] if FA_present else np.zeros(0, dtype=int))
1✔
2974

2975
        # boolean index vector into SNR
2976
        #   True iff we have computed a SNR for that planet
2977
        #   False iff we didn't find an SNR (and will leave 0 there)
2978
        #   if FA_present, SNR_plug_in[-1] = True
2979
        SNR_plug_in = np.isin(pInds_FA, pIndsChar)
1✔
2980

2981
        # save all SNRs (planets and FA) to one array
2982
        SNR = np.zeros(len(pInds_FA))
1✔
2983
        # plug in the SNR's we computed (pIndsChar and optional FA)
2984
        SNR[SNR_plug_in] = np.append(SNRplans, SNRfa)
1✔
2985

2986
        return SNR, fZ, systemParams
1✔
2987

2988
    def revisit_inds(self, sInds, tmpCurrentTimeNorm):
1✔
2989
        """Return subset of star indices that are scheduled for revisit.
2990

2991
        Args:
2992
            sInds (~numpy.ndarray(int)):
2993
                Indices of stars to consider
2994
            tmpCurrentTimeNorm (astropy.units.Quantity):
2995
                Normalized simulation time
2996

2997
        Returns:
2998
            ~numpy.ndarray(int):
2999
                Indices of stars whose revisit is scheduled for within `self.dt_max` of
3000
                the current time
3001

3002
        """
3003
        dt_rev = np.abs(self.starRevisit[:, 1] * u.day - tmpCurrentTimeNorm)
1✔
3004
        ind_rev = [
1✔
3005
            int(x) for x in self.starRevisit[dt_rev < self.dt_max, 0] if x in sInds
3006
        ]
3007

3008
        return ind_rev
1✔
3009

3010
    def keepout_filter(self, sInds, startTimes, endTimes, koMap):
1✔
3011
        """Filters stars not observable due to keepout
3012

3013
        Args:
3014
            sInds (~numpy.ndarray(int)):
3015
                indices of stars still in observation list
3016
            startTimes (~numpy.ndarray(float)):
3017
                start times of observations
3018
            endTimes (~numpy.ndarray(float)):
3019
                end times of observations
3020
            koMap (~numpy.ndarray(bool)):
3021
                map where True is for a target unobstructed and observable,
3022
                False is for a target obstructed and unobservable due to keepout zone
3023

3024
        Returns:
3025
            ~numpy.ndarray(int):
3026
                sInds - filtered indices of stars still in observation list
3027

3028
        """
3029
        # find the indices of keepout times that pertain to the start and end times of
3030
        # observation
3031
        startInds = np.searchsorted(self.koTimes.value, startTimes)
×
3032
        endInds = np.searchsorted(self.koTimes.value, endTimes)
×
3033

3034
        # boolean array of available targets (unobstructed between observation start and
3035
        # end time) we include a -1 in the start and +1 in the end to include keepout
3036
        # days enclosing start and end times
3037
        availableTargets = np.array(
×
3038
            [
3039
                np.all(
3040
                    koMap[
3041
                        sInd,
3042
                        max(startInds[sInd] - 1, 0) : min(
3043
                            endInds[sInd] + 1, len(self.koTimes.value) + 1
3044
                        ),
3045
                    ]
3046
                )
3047
                for sInd in sInds
3048
            ],
3049
            dtype="bool",
3050
        )
3051

3052
        return sInds[availableTargets]
×
3053

3054

3055
def array_encoder(obj):
1✔
3056
    r"""Encodes numpy arrays, astropy Times, and astropy Quantities, into JSON.
3057

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

3064
    Args:
3065
        obj (Any):
3066
            Object to encode.
3067

3068
        Returns:
3069
            Any:
3070
                Encoded object
3071

3072
    """
3073

3074
    from astropy.coordinates import SkyCoord
1✔
3075
    from astropy.time import Time
1✔
3076

3077
    if isinstance(obj, Time):
1✔
3078
        # astropy Time -> time string
3079
        return obj.fits  # isot also makes sense here
×
3080
    if isinstance(obj, u.quantity.Quantity):
1✔
3081
        # note: it is possible to have a numpy ndarray wrapped in a Quantity.
3082
        # NB: alternatively, can return (obj.value, obj.unit.name)
3083
        return obj.value
×
3084
    if isinstance(obj, SkyCoord):
1✔
3085
        return dict(
×
3086
            lon=obj.heliocentrictrueecliptic.lon.value,
3087
            lat=obj.heliocentrictrueecliptic.lat.value,
3088
            distance=obj.heliocentrictrueecliptic.distance.value,
3089
        )
3090
    if isinstance(obj, (np.ndarray, np.number)):
1✔
3091
        # ndarray -> list of numbers
3092
        return obj.tolist()
1✔
3093
    if isinstance(obj, complex):
×
3094
        # complex -> (real, imag) pair
3095
        return [obj.real, obj.imag]
×
3096
    if callable(obj):
×
3097
        # this case occurs for interpolants like PSF and QE
3098
        # We cannot simply "write" the function to JSON, so we make up a string
3099
        # to keep from throwing an error.
3100
        # The fix is simple: when generating the interpolant, add a _outspec attribute
3101
        # to the function (or the lambda), containing (e.g.) the fits filename, or the
3102
        # explicit number -- whatever string was used.  Then, here, check for that
3103
        # attribute and write it out instead of this dummy string.  (Attributes can
3104
        # be transparently attached to python functions, even lambda's.)
3105
        return "interpolant_function"
×
3106
    if isinstance(obj, set):
×
3107
        return list(obj)
×
3108
    if isinstance(obj, bytes):
×
3109
        return obj.decode()
×
3110
    # an EXOSIMS object
3111
    if hasattr(obj, "_modtype"):
×
3112
        return obj.__dict__
×
3113
    # an object for which no encoding is defined yet
3114
    #   as noted above, ordinary types (lists, ints, floats) do not take this path
3115
    raise ValueError("Could not JSON-encode an object of type %s" % type(obj))
×
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE TRIAL · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc