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

dsavransky / EXOSIMS / 14561753842

20 Apr 2025 05:25PM UTC coverage: 65.777% (+0.07%) from 65.708%
14561753842

push

github

web-flow
Merge pull request #408 from dsavransky/performance_optimization

Performance optimizations.

604 of 731 new or added lines in 15 files covered. (82.63%)

22 existing lines in 9 files now uncovered.

9681 of 14718 relevant lines covered (65.78%)

0.66 hits per line

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

62.37
/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 (only if ``keepStarCatalog`` input is True.
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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

529
    def initializeStorageArrays(self):
1✔
530
        """
531
        Initialize all storage arrays based on # of stars and targets
532
        """
533

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

546
    def __str__(self):
1✔
547
        """String representation of the Survey Simulation object
548

549
        When the command 'print' is used on the Survey Simulation object, this
550
        method will return the values contained in the object
551

552
        """
553

554
        for att in self.__dict__:
1✔
555
            print("%s: %r" % (att, getattr(self, att)))
1✔
556

557
        return "Survey Simulation class object attributes"
1✔
558

559
    def run_sim(self):
1✔
560
        """Performs the survey simulation"""
561

562
        OS = self.OpticalSystem
1✔
563
        TL = self.TargetList
1✔
564
        SU = self.SimulatedUniverse
1✔
565
        Obs = self.Observatory
1✔
566
        TK = self.TimeKeeping
1✔
567

568
        # TODO: start using this self.currentSep
569
        # set occulter separation if haveOcculter
570
        if OS.haveOcculter:
1✔
571
            self.currentSep = Obs.occulterSep
×
572

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

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

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

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

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

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

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

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

699
                # append result values to self.DRM
700
                self.DRM.append(DRM)
1✔
701

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

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

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

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

821
        Returns:
822
            None
823
        """
824

825
        self.TimeKeeping.allocate_time(
1✔
826
            dt
827
            * (1.0 - self.TimeKeeping.missionPortion)
828
            / self.TimeKeeping.missionPortion,
829
            addExoplanetObsTime=False,
830
        )
831

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

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

839
        Args:
840
            old_sInd (int):
841
                Index of the previous target star
842
            mode (dict):
843
                Selected observing mode for detection
844

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

859
        """
860
        OS = self.OpticalSystem
1✔
861
        TL = self.TargetList
1✔
862
        Obs = self.Observatory
1✔
863
        TK = self.TimeKeeping
1✔
864

865
        # create DRM
866
        DRM = {}
1✔
867

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

876
        # create appropriate koMap
877
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
878

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

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

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

904
        # start times, including slew times
905
        startTimes = tmpCurrentTimeAbs.copy() + slewTimes
1✔
906
        startTimesNorm = tmpCurrentTimeNorm.copy() + slewTimes
1✔
907

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

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

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

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

957
                if maxIntTime.value <= 0:
1✔
958
                    sInds = np.asarray([], dtype=int)
×
959

960
        # 5.1 TODO Add filter to filter out stars entering and exiting keepout
961
        # between startTimes and endTimes
962

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

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

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

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

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

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

1027
        return DRM, sInd, intTime, waitTime
1✔
1028

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

1033
        Given a subset of targets, calculate their integration times given the
1034
        start of observation time.
1035

1036
        Prototype just calculates integration times for fixed contrast depth.
1037

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

1047
        Returns:
1048
            ~astropy.units.Quantity(~numpy.ndarray(float)):
1049
                Integration times for detection
1050
                same dimension as sInds
1051

1052
        .. note::
1053

1054
            next_target filter will discard targets with zero integration times.
1055

1056
        """
1057

1058
        TL = self.TargetList
1✔
1059

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

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

1076
            count_fpath = os.path.join(self.record_counts_path, "counts")
×
1077

1078
            if not os.path.exists(count_fpath):
×
1079
                os.mkdir(count_fpath)
×
1080

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

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

1118
            with open(outfile, "w") as csvfile:
×
1119
                c = csv.writer(csvfile)
×
1120
                c.writerows(self.count_lines)
×
1121

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

1127
        return intTimes
1✔
1128

1129
    def choose_next_target(self, old_sInd, sInds, slewTimes, intTimes):
1✔
1130
        """Helper method for method next_target to simplify alternative implementations.
1131

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

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

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

1154
        """
1155

1156
        Comp = self.Completeness
1✔
1157
        TL = self.TargetList
1✔
1158
        TK = self.TimeKeeping
1✔
1159
        OS = self.OpticalSystem
1✔
1160
        Obs = self.Observatory
1✔
1161
        allModes = OS.observingModes
1✔
1162

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

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

1193
        return (
1✔
1194
            sInd,
1195
            slewTimes[sInd],
1196
        )  # if coronagraph or first sInd, waitTime will be 0 days
1197

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

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

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

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

1237
        Obs = self.Observatory
×
1238
        TL = self.TargetList
×
1239

1240
        # initializing arrays
NEW
1241
        obsTimeArray = np.zeros([TL.nStars, 50]) << u.d
×
NEW
1242
        intTimeArray = np.zeros([TL.nStars, 2]) << u.d
×
1243

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

1255
        # determining which scheme to use to filter slews
1256
        obsModName = Obs.__class__.__name__
×
1257

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

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

1281
        return sInds, slewTimes, intTimes, dV
×
1282

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

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

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

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

1315
        TK = self.TimeKeeping
×
1316
        Obs = self.Observatory
×
1317

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1437
            slewTimes = slewTimes[good_inds]
×
1438
        else:
1439
            slewTimes = slewTimes[good_inds]
×
1440

1441
        return sInds[good_inds], intTimes[good_inds].flatten(), slewTimes
×
1442

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

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

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

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

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

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

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

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

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

1541
        maxObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1542
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1543

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

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

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

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

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

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

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

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

1602
        if any(OB_withObsStars > 0):
×
1603
            nOBstart = np.argmin(np.abs(OB_withObsStars))
×
1604
            nOBend = np.argmax(OB_withObsStars)
×
1605

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

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

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

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

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

1658
                minAllowedSlewTimes_nOB[np.invert(conds)] = np.inf
×
1659
                maxAllowedSlewTimes_nOB[np.invert(conds)] = -np.inf
×
1660

1661
                # one last condition
1662
                map_i, map_j = np.where(
×
1663
                    (obsTimeArrayNorm > minAllowedSlewTimes_nOB)
1664
                    & (obsTimeArrayNorm < maxAllowedSlewTimes_nOB)
1665
                )
1666

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

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

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

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

1704
            return sInds, intTime, slewTime, dV
×
1705

1706
        else:
1707
            empty = np.asarray([], dtype=int)
×
NEW
1708
            return empty, empty << u.d, empty << u.d, empty << self.m_per_s
×
1709

1710
    def chooseOcculterSlewTimes(self, sInds, slewTimes, dV, intTimes, charTimes):
1✔
1711
        """Selects the best slew time for each star
1712

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

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

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

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

1749
        dV = dV[good_i, good_j]
×
1750
        intTime = intTimes[good_i, good_j]
×
1751
        slewTime = slewTimes[good_i, good_j]
×
1752

1753
        return sInds, intTime, slewTime, dV
×
1754

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

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

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

1785
        """
1786

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

1795
        # Save Current Time before attempting time allocation
1796
        currentTimeNorm = TK.currentTimeNorm.copy()
1✔
1797
        currentTimeAbs = TK.currentTimeAbs.copy()
1✔
1798

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

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

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

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

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

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

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

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

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

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

1953
        # Schedule Target Revisit
1954
        self.scheduleRevisit(sInd, smin, det, pInds)
1✔
1955

1956
        if self.make_debug_bird_plots:
1✔
1957
            from tools.obs_plot import obs_plot
×
1958

1959
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, detected)
×
1960

1961
        return detected.astype(int), fZ, JEZ, systemParams, SNR, FA
1✔
1962

1963
    def scheduleRevisit(self, sInd, smin, det, pInds):
1✔
1964
        """A Helper Method for scheduling revisits after observation detection
1965

1966
        Updates self.starRevisit attribute only
1967

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

1978
        Returns:
1979
            None
1980

1981
        """
1982
        TK = self.TimeKeeping
1✔
1983
        TL = self.TargetList
1✔
1984
        SU = self.SimulatedUniverse
1✔
1985

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

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

2022
    def observation_characterization(self, sInd, mode):
1✔
2023
        """Finds if characterizations are possible and relevant information
2024

2025
        Args:
2026
            sInd (int):
2027
                Integer index of the star of interest
2028
            mode (dict):
2029
                Selected observing mode for characterization
2030

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

2051
        """
2052

2053
        OS = self.OpticalSystem
1✔
2054
        ZL = self.ZodiacalLight
1✔
2055
        TL = self.TargetList
1✔
2056
        SU = self.SimulatedUniverse
1✔
2057
        Obs = self.Observatory
1✔
2058
        TK = self.TimeKeeping
1✔
2059

2060
        # selecting appropriate koMap
2061
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
2062

2063
        # find indices of planets around the target
2064
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2065

2066
        # get the detected status, and check if there was a FA
2067
        det = self.lastDetected[sInd, 0]
1✔
2068
        FA = len(det) == len(pInds) + 1
1✔
2069
        if FA:
1✔
2070
            pIndsDet = np.append(pInds, -1)[det]
×
2071
        else:
2072
            pIndsDet = pInds[det]
1✔
2073

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

2086
        # look for last detected planets that have not been fully characterized
2087
        if not (FA):  # only true planets, no FA
1✔
2088
            tochar = self.fullSpectra[pIndsDet] == 0
1✔
2089
        else:  # mix of planets and a FA
2090
            truePlans = pIndsDet[:-1]
×
2091
            tochar = np.append((self.fullSpectra[truePlans] == 0), True)
×
2092

2093
        # 1/ find spacecraft orbital START position including overhead time,
2094
        # and check keepout angle
2095
        if np.any(tochar):
1✔
2096
            # start times
2097
            startTime = (
1✔
2098
                TK.currentTimeAbs.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2099
            )
2100

2101
            # if we're beyond mission end, bail out
2102
            if startTime >= TK.missionFinishAbs:
1✔
2103
                return characterized, fZ, systemParams, SNR, intTime
×
2104

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

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

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

2138
            # end times
2139
            endTimes = startTime + totTimes
1✔
2140
            endTimesNorm = startTimeNorm + totTimes
1✔
2141
            # planets to characterize
2142
            tochar = (
1✔
2143
                (totTimes > 0)
2144
                & (totTimes <= OS.intCutoff)
2145
                & (endTimesNorm <= TK.OBendTimes[TK.OBnumber])
2146
            )
2147

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

2169
        # 4/ if yes, allocate the overhead time, and perform the characterization
2170
        # for the maximum char time
2171
        if np.any(tochar):
1✔
2172
            # Save Current Time before attempting time allocation
2173
            currentTimeNorm = TK.currentTimeNorm.copy()
1✔
2174
            currentTimeAbs = TK.currentTimeAbs.copy()
1✔
2175

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

2199
            pIndsChar = pIndsDet[tochar]
1✔
2200
            log_char = "   - Charact. planet inds %s (%s/%s detected)" % (
1✔
2201
                pIndsChar,
2202
                len(pIndsChar),
2203
                len(pIndsDet),
2204
            )
2205
            self.logger.info(log_char)
1✔
2206
            self.vprint(log_char)
1✔
2207

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

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

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

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

2290
            # save all SNRs (planets and FA) to one array
2291
            SNRinds = np.where(det)[0][tochar]
1✔
2292
            SNR[SNRinds] = np.append(SNRplans, SNRfa)
1✔
2293

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

2315
        if self.make_debug_bird_plots:
1✔
2316
            from tools.obs_plot import obs_plot
×
2317

2318
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2319

2320
        return characterized.astype(int), fZ, JEZ, systemParams, SNR, intTime
1✔
2321

2322
    def calc_signal_noise(
1✔
2323
        self, sInd, pInds, t_int, mode, fZ=None, JEZ=None, dMag=None, WA=None
2324
    ):
2325
        """Calculates the signal and noise fluxes for a given time interval. Called
2326
        by observation_detection and observation_characterization methods in the
2327
        SurveySimulation module.
2328

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

2347
        Returns:
2348
            tuple:
2349
                Signal (float):
2350
                    Counts of signal
2351
                Noise (float):
2352
                    Counts of background noise variance
2353

2354
        """
2355

2356
        OS = self.OpticalSystem
1✔
2357
        ZL = self.ZodiacalLight
1✔
2358
        TL = self.TargetList
1✔
2359
        SU = self.SimulatedUniverse
1✔
2360
        Obs = self.Observatory
1✔
2361
        TK = self.TimeKeeping
1✔
2362

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

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

2401
        # initialize Signal and Noise arrays
2402
        Signal = np.zeros(len(pInds))
1✔
2403
        Noise = np.zeros(len(pInds))
1✔
2404

2405
        # find observable planets wrt IWA-OWA range
2406
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2407

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

2416
            Signal[obs] = C_p.to_value(self.inv_s) * _t_int_s
1✔
2417
            Noise[obs] = np.sqrt(
1✔
2418
                (
2419
                    C_b.to_value(self.inv_s) * _t_int_s
2420
                    + (C_sp.to_value(self.inv_s) * _t_int_s) ** 2
2421
                )
2422
            )
2423

2424
        return Signal, Noise
1✔
2425

2426
    def update_occulter_mass(self, DRM, sInd, t_int, skMode):
1✔
2427
        """Updates the occulter wet mass in the Observatory module, and stores all
2428
        the occulter related values in the DRM array.
2429

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

2442
        Returns:
2443
            dict:
2444
                Design Reference Mission dictionary, contains the results of one
2445
                complete observation
2446

2447
        """
2448

2449
        TL = self.TargetList
×
2450
        Obs = self.Observatory
×
2451
        TK = self.TimeKeeping
×
2452

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

2455
        # decrement mass for station-keeping
2456
        dF_lateral, dF_axial, intMdot, mass_used, deltaV = Obs.mass_dec_sk(
×
2457
            TL, sInd, TK.currentTimeAbs.copy(), t_int
2458
        )
2459

2460
        DRM[skMode + "_dV"] = deltaV.to("m/s")
×
2461
        DRM[skMode + "_mass_used"] = mass_used.to("kg")
×
2462
        DRM[skMode + "_dF_lateral"] = dF_lateral.to("N")
×
2463
        DRM[skMode + "_dF_axial"] = dF_axial.to("N")
×
2464
        # update spacecraft mass
2465
        Obs.scMass = Obs.scMass - mass_used
×
2466
        DRM["scMass"] = Obs.scMass.to("kg")
×
2467
        if Obs.twotanks:
×
2468
            Obs.skMass = Obs.skMass - mass_used
×
2469
            DRM["skMass"] = Obs.skMass.to("kg")
×
2470

2471
        return DRM
×
2472

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

2476
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2477
        objects with their own outspecs.
2478

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

2496
        """
2497

2498
        SU = self.SimulatedUniverse
1✔
2499
        TK = self.TimeKeeping
1✔
2500
        TL = self.TargetList
1✔
2501

2502
        # re-initialize SurveySimulation arrays
2503
        specs = self._outspec
1✔
2504
        specs["modules"] = self.modules
1✔
2505

2506
        if seed is None:  # pop the seed so a new one is generated
1✔
2507
            if "seed" in specs:
1✔
2508
                specs.pop("seed")
1✔
2509
        else:  # if seed is provided, replace seed with provided seed
2510
            specs["seed"] = seed
×
2511

2512
        # reset mission time, re-init surveysim and observatory
2513
        TK.__init__(**TK._outspec)
1✔
2514
        self.__init__(**specs)
1✔
2515
        self.Observatory.__init__(**self.Observatory._outspec)
1✔
2516

2517
        # generate new planets if requested (default)
2518
        if genNewPlanets:
1✔
2519
            TL.stellar_mass()
1✔
2520
            TL.I = TL.gen_inclinations(TL.PlanetPopulation.Irange)  # noqa: E741
1✔
2521
            SU.gen_physical_properties(**SU._outspec)
1✔
2522
            rewindPlanets = True
1✔
2523
        # re-initialize systems if requested (default)
2524
        if rewindPlanets:
1✔
2525
            SU.init_systems()
1✔
2526

2527
        # reset helper arrays
2528
        self.initializeStorageArrays()
1✔
2529

2530
        self.vprint("Simulation reset.")
1✔
2531

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

2541
        Args:
2542
            tofile (str, optional):
2543
                Name of the file containing all output specifications (outspecs).
2544
                Defaults to None.
2545
            starting_outspec (dict, optional):
2546
                Initial outspec (from MissionSim). Defaults to None.
2547
            modnames (bool):
2548
                If True, populate outspec dictionary with the module it originated from,
2549
                instead of the actual value of the keyword. Defaults False.
2550

2551
        Returns:
2552
            dict:
2553
                Dictionary containing the full :ref:`sec:inputspec`, including all
2554
                filled-in default values. Combination of all individual module _outspec
2555
                attributes.
2556

2557
        """
2558

2559
        # start with our own outspec
2560
        if modnames:
1✔
2561
            out = copy.copy(self._outspec)
1✔
2562
            for k in out:
1✔
2563
                out[k] = self.__class__.__name__
1✔
2564
        else:
2565
            out = copy.deepcopy(self._outspec)
1✔
2566

2567
        # Add any provided other outspec
2568
        if starting_outspec is not None:
1✔
2569
            out.update(starting_outspec)
1✔
2570

2571
        # add in all modules _outspec's
2572
        for module in self.modules.values():
1✔
2573
            if modnames:
1✔
2574
                tmp = copy.copy(module._outspec)
1✔
2575
                for k in tmp:
1✔
2576
                    tmp[k] = module.__class__.__name__
1✔
2577
            else:
2578
                tmp = module._outspec
1✔
2579
            out.update(tmp)
1✔
2580

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

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

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

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

2633
        return out
1✔
2634

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

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

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

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

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

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

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

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

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

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

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

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

2737

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2867
        """
2868

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

2876
        # time at start of integration window
2877
        currentTimeNorm = startTime
×
2878
        currentTimeAbs = TK.missionStart + startTime
×
2879

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

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

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

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

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

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

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

2987
        return SNR, fZ, systemParams
×
2988

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

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

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

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

3009
        return ind_rev
1✔
3010

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

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

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

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

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

3053
        return sInds[availableTargets]
×
3054

3055

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

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

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

3069
        Returns:
3070
            Any:
3071
                Encoded object
3072

3073
    """
3074

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

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

© 2026 Coveralls, Inc