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

dsavransky / EXOSIMS / 24738743655

21 Apr 2026 06:12PM UTC coverage: 65.925% (-0.09%) from 66.015%
24738743655

Pull #457

github

web-flow
Merge 894d13161 into e04cd29e1
Pull Request #457: YAML Input Script Capability

29 of 36 new or added lines in 2 files covered. (80.56%)

20 existing lines in 5 files now uncovered.

9896 of 15011 relevant lines covered (65.92%)

0.66 hits per line

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

66.35
/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 yaml
1✔
7
import logging
1✔
8
import os
1✔
9
import random as py_random
1✔
10
import re
1✔
11
import time
1✔
12
from pathlib import Path
1✔
13
from typing import Any, Dict, Optional
1✔
14

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

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

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

29

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

33
    Args:
34

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

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

188
    """
189

190
    _modtype = "SurveySimulation"
1✔
191

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

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

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

233
            if scriptfile.lower().endswith(".yaml") or scriptfile.lower().endswith(
1✔
234
                ".yml"
235
            ):
NEW
236
                try:
×
NEW
237
                    specs.update(yaml.safe_load(script))
×
NEW
238
                except ValueError:
×
NEW
239
                    print(
×
240
                        "Error: %s: Script file `%s' is not valid YAML."
241
                        % (self._modtype, scriptfile)
242
                    )
243
                    # must re-raise, or the error will be masked
NEW
244
                    raise
×
245
            elif scriptfile.lower().endswith(".json"):
1✔
246
                try:
1✔
247
                    specs.update(json.loads(script))
1✔
248
                except ValueError:
1✔
249
                    print(
1✔
250
                        "Error: %s: Script file `%s' is not valid JSON."
251
                        % (self._modtype, scriptfile)
252
                    )
253
                    # must re-raise, or the error will be masked
254
                    raise
1✔
255
            else:  # Force file extension to be either .json or .yaml/.yml
NEW
256
                raise ValueError(
×
257
                    "Error: %s: Input file `%s' file extension not JSON or YAML."
258
                    % (self._modtype, scriptfile)
259
                )
260

261
            # modules array must be present
262
            if "modules" not in specs:
1✔
263
                raise ValueError("No modules field found in script.")
×
264

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

268
        # count dict contains all of the C info for each star index
269
        self.record_counts_path = record_counts_path
1✔
270
        self.count_lines = []
1✔
271
        self._outspec["record_counts_path"] = record_counts_path
1✔
272

273
        # mission simulation logger
274
        self.logger = specs.get("logger", logging.getLogger(__name__))
1✔
275

276
        # set up numpy random number (generate it if not in specs)
277
        self.seed = int(specs.get("seed", py_random.randint(1, int(1e9))))
1✔
278
        self.vprint("Numpy random seed is: %s" % self.seed)
1✔
279
        np.random.seed(self.seed)
1✔
280
        self._outspec["seed"] = self.seed
1✔
281

282
        # cache directory
283
        self.cachedir = get_cache_dir(cachedir)
1✔
284
        self._outspec["cachedir"] = self.cachedir
1✔
285
        # N.B.: cachedir is going to be used by everything, so let's make sure that
286
        # it doesn't get popped out of specs
287
        specs["cachedir"] = self.cachedir
1✔
288

289
        # if any of the modules is a string, assume that they are all strings
290
        # and we need to initalize
291
        if isinstance(next(iter(specs["modules"].values())), str):
1✔
292
            # import desired module names (prototype or specific)
293
            self.SimulatedUniverse = get_module(
1✔
294
                specs["modules"]["SimulatedUniverse"], "SimulatedUniverse"
295
            )(**specs)
296
            self.Observatory = get_module(
1✔
297
                specs["modules"]["Observatory"], "Observatory"
298
            )(**specs)
299
            self.TimeKeeping = get_module(
1✔
300
                specs["modules"]["TimeKeeping"], "TimeKeeping"
301
            )(**specs)
302

303
            # bring inherited class objects to top level of Survey Simulation
304
            SU = self.SimulatedUniverse
1✔
305
            self.StarCatalog = SU.StarCatalog
1✔
306
            self.PlanetPopulation = SU.PlanetPopulation
1✔
307
            self.PlanetPhysicalModel = SU.PlanetPhysicalModel
1✔
308
            self.OpticalSystem = SU.OpticalSystem
1✔
309
            self.ZodiacalLight = SU.ZodiacalLight
1✔
310
            self.BackgroundSources = SU.BackgroundSources
1✔
311
            self.PostProcessing = SU.PostProcessing
1✔
312
            self.Completeness = SU.Completeness
1✔
313
            self.TargetList = SU.TargetList
1✔
314

315
        else:
316
            # these are the modules that must be present if passing instantiated objects
317
            neededObjMods = [
1✔
318
                "PlanetPopulation",
319
                "PlanetPhysicalModel",
320
                "OpticalSystem",
321
                "ZodiacalLight",
322
                "BackgroundSources",
323
                "PostProcessing",
324
                "Completeness",
325
                "TargetList",
326
                "SimulatedUniverse",
327
                "Observatory",
328
                "TimeKeeping",
329
            ]
330

331
            # ensure that you have the minimal set
332
            for modName in neededObjMods:
1✔
333
                if modName not in specs["modules"]:
1✔
334
                    raise ValueError(
×
335
                        "%s module is required but was not provided." % modName
336
                    )
337

338
            for modName in specs["modules"]:
1✔
339
                assert specs["modules"][modName]._modtype == modName, (
1✔
340
                    "Provided instance of %s has incorrect modtype." % modName
341
                )
342

343
                setattr(self, modName, specs["modules"][modName])
1✔
344

345
        # create a dictionary of all modules, except StarCatalog
346
        self.modules = {}
1✔
347
        self.modules["PlanetPopulation"] = self.PlanetPopulation
1✔
348
        self.modules["PlanetPhysicalModel"] = self.PlanetPhysicalModel
1✔
349
        self.modules["OpticalSystem"] = self.OpticalSystem
1✔
350
        self.modules["ZodiacalLight"] = self.ZodiacalLight
1✔
351
        self.modules["BackgroundSources"] = self.BackgroundSources
1✔
352
        self.modules["PostProcessing"] = self.PostProcessing
1✔
353
        self.modules["Completeness"] = self.Completeness
1✔
354
        self.modules["TargetList"] = self.TargetList
1✔
355
        self.modules["SimulatedUniverse"] = self.SimulatedUniverse
1✔
356
        self.modules["Observatory"] = self.Observatory
1✔
357
        self.modules["TimeKeeping"] = self.TimeKeeping
1✔
358
        # add yourself to modules list for bookkeeping purposes
359
        self.modules["SurveySimulation"] = self
1✔
360

361
        # observation time sampling
362
        self.ntFlux = int(ntFlux)
1✔
363
        self._outspec["ntFlux"] = self.ntFlux
1✔
364

365
        # maximum number of observations per star
366
        self.nVisitsMax = int(nVisitsMax)
1✔
367
        self._outspec["nVisitsMax"] = self.nVisitsMax
1✔
368

369
        # integration time margin for characterization
370
        self.charMargin = float(charMargin)
1✔
371
        self._outspec["charMargin"] = self.charMargin
1✔
372

373
        # maximum time for revisit window
374
        self.dt_max = float(dt_max) * u.week
1✔
375
        self._outspec["dt_max"] = self.dt_max.value
1✔
376

377
        # minimum time for revisit window
378
        if revisit_wait is not None:
1✔
379
            self.revisit_wait = revisit_wait * u.d
×
380
        else:
381
            self.revisit_wait = revisit_wait
1✔
382
        self._outspec["revisit_wait"] = revisit_wait
1✔
383

384
        # list of detected earth-like planets aroung promoted stars
385
        self.known_earths = np.array([])
1✔
386

387
        self.find_known_RV = find_known_RV
1✔
388
        self._outspec["find_known_RV"] = find_known_RV
1✔
389
        self._outspec["include_known_RV"] = include_known_RV
1✔
390
        if self.find_known_RV:
1✔
391
            # select specific knonw RV stars if a file exists
392
            if include_known_RV is not None:
×
393
                if os.path.isfile(include_known_RV):
×
394
                    with open(include_known_RV, "r") as rv_file:
×
395
                        self.include_known_RV = [
×
396
                            hip.strip() for hip in rv_file.read().split(",")
397
                        ]
398
                        self.vprint(
×
399
                            "Including known RV stars: {}".format(self.include_known_RV)
400
                        )
401
                else:
402
                    self.include_known_RV = None
×
403
                    self.vprint(
×
404
                        "WARNING: Known RV file: {} does not exist!!".format(
405
                            include_known_RV
406
                        )
407
                    )
408
            else:
409
                self.include_known_RV = None
×
410
            self.known_stars, self.known_rocky = self.find_known_plans()
×
411
        else:
412
            self.include_known_RV = None
1✔
413
            self.known_stars = []
1✔
414
            self.known_rocky = []
1✔
415

416
        # defaultAddExoplanetObsTime Tells us time advanced when no targets available
417
        # counts agains exoplanetObsTime (when True)
418
        self.defaultAddExoplanetObsTime = defaultAddExoplanetObsTime
1✔
419
        self._outspec["defaultAddExoplanetObsTime"] = defaultAddExoplanetObsTime
1✔
420

421
        # If inputs are scalars, save scalars to outspec, otherwise save full lists
422
        OS = self.OpticalSystem
1✔
423
        TL = self.TargetList
1✔
424
        SU = self.SimulatedUniverse
1✔
425
        TK = self.TimeKeeping
1✔
426

427
        # initialize arrays updated in run_sim()
428
        self.initializeStorageArrays()
1✔
429

430
        # Generate File Hashnames and loction
431
        self.cachefname = self.generateHashfName(specs)
1✔
432

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

437
        # getting keepout map for entire mission
438
        startTime = self.TimeKeeping.missionStart.copy()
1✔
439
        endTime = self.TimeKeeping.missionFinishAbs.copy()
1✔
440

441
        nSystems = len(allModes)
1✔
442
        systNames = np.unique(
1✔
443
            [allModes[x]["syst"]["name"] for x in np.arange(nSystems)]
444
        )
445
        systOrder = np.argsort(systNames)
1✔
446
        koStr = ["koAngles_Sun", "koAngles_Moon", "koAngles_Earth", "koAngles_Small"]
1✔
447
        koangles = np.zeros([len(systNames), 4, 2])
1✔
448

449
        for x in systOrder:
1✔
450
            rel_mode = list(
1✔
451
                filter(lambda mode: mode["syst"]["name"] == systNames[x], allModes)
452
            )[0]
453
            koangles[x] = np.asarray([rel_mode["syst"][k] for k in koStr])
1✔
454

455
        # Precalculate Earth positions and create a cubic spline interpolator
456
        # to save time during orbit calculations
457
        dt = 0.1  # days
1✔
458
        self.Observatory.generate_Earth_interpolator(startTime, endTime, dt)
1✔
459

460
        self._outspec["nokoMap"] = nokoMap
1✔
461
        if not (nokoMap):
1✔
462
            koMaps, self.koTimes = self.Observatory.generate_koMap(
1✔
463
                TL, startTime, endTime, koangles
464
            )
465
            self.koTimes_mjd = self.koTimes.mjd
1✔
466
            self.koMaps = {}
1✔
467
            for x, n in zip(systOrder, systNames[systOrder]):
1✔
468
                print(n)
1✔
469
                self.koMaps[n] = koMaps[x, :, :]
1✔
470

471
        self._outspec["nofZ"] = nofZ
1✔
472

473
        # TODO: do we still want a nofZ option?  probably not.
474
        self.fZmins = {}
1✔
475
        self.fZtypes = {}
1✔
476
        for x, n in zip(systOrder, systNames[systOrder]):
1✔
477
            self.fZmins[n] = np.array([])
1✔
478
            self.fZtypes[n] = np.array([])
1✔
479

480
        # TODO: do we need to do this for all modes? doing det only breaks other
481
        # schedulers, but maybe there's a better approach here.
482
        sInds = np.arange(TL.nStars)  # Initialize some sInds array
1✔
483
        for mode in allModes:
1✔
484
            # This instantiates fZMap arrays for every starlight suppresion system
485
            # that is actually used in a mode
486
            modeHashName = (
1✔
487
                f'{self.cachefname[0:-1]}_{TL.nStars}_{mode["syst"]["name"]}'
488
                f"_{startTime.mjd:0}_{endTime.mjd:0}."
489
            )
490

491
            if (np.size(self.fZmins[mode["syst"]["name"]]) == 0) or (
1✔
492
                np.size(self.fZtypes[mode["syst"]["name"]]) == 0
493
            ):
494
                self.ZodiacalLight.generate_fZ(
1✔
495
                    self.Observatory,
496
                    TL,
497
                    self.TimeKeeping,
498
                    mode,
499
                    modeHashName,
500
                    self.koTimes,
501
                )
502

503
                (
1✔
504
                    self.fZmins[mode["syst"]["name"]],
505
                    self.fZtypes[mode["syst"]["name"]],
506
                ) = self.ZodiacalLight.calcfZmin(
507
                    sInds,
508
                    self.Observatory,
509
                    TL,
510
                    self.TimeKeeping,
511
                    mode,
512
                    modeHashName,
513
                    self.koMaps[mode["syst"]["name"]],
514
                    self.koTimes,
515
                )
516

517
        # Precalculating intTimeFilter for coronagraph
518
        # find fZmin to use in intTimeFilter
519
        self.valfZmin, self.absTimefZmin = self.ZodiacalLight.extractfZmin(
1✔
520
            self.fZmins[det_mode["syst"]["name"]], sInds, self.koTimes
521
        )
522
        JEZ0 = self.TargetList.JEZ0[det_mode["hex"]]
1✔
523
        dMag = TL.int_dMag[sInds]  # grabbing dMag
1✔
524
        WA = TL.int_WA[sInds]  # grabbing WA
1✔
525
        self.intTimesIntTimeFilter = (
1✔
526
            self.OpticalSystem.calc_intTime(
527
                TL, sInds, self.valfZmin, JEZ0, dMag, WA, det_mode, TK=TK
528
            )
529
            * det_mode["timeMultiplier"]
530
        )  # intTimes to filter by
531
        # These indices are acceptable for use simulating
532
        self.intTimeFilterInds = np.where(
1✔
533
            (
534
                (self.intTimesIntTimeFilter > 0)
535
                & (self.intTimesIntTimeFilter <= self.OpticalSystem.intCutoff)
536
            )
537
        )[0]
538

539
        self.make_debug_bird_plots = make_debug_bird_plots
1✔
540
        if self.make_debug_bird_plots:
1✔
541
            assert (
×
542
                debug_plot_path is not None
543
            ), "debug_plot_path must be set by input if make_debug_bird_plots is True"
544
            self.obs_plot_path = Path(f"{debug_plot_path}/{self.seed}")
×
545
            # Make directory if it doesn't exist
546
            if not self.obs_plot_path.exists():
×
547
                vprint(f"Making plot directory: {self.obs_plot_path}")
×
548
                self.obs_plot_path.mkdir(parents=True, exist_ok=True)
×
549
            self.obs_n_counter = 0
×
550
        self._outspec["make_debug_bird_plots"] = self.make_debug_bird_plots
1✔
551
        self._outspec["debug_plot_path"] = debug_plot_path
1✔
552

553
    def initializeStorageArrays(self):
1✔
554
        """
555
        Initialize all storage arrays based on # of stars and targets
556
        """
557

558
        self.DRM = []
1✔
559
        self.fullSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
1✔
560
        self.partialSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
1✔
561
        self.propagTimes = np.zeros(self.TargetList.nStars) * u.d
1✔
562
        self.lastObsTimes = np.zeros(self.TargetList.nStars) * u.d
1✔
563
        self.starVisits = np.zeros(
1✔
564
            self.TargetList.nStars, dtype=int
565
        )  # contains the number of times each star was visited
566
        self.starRevisit = np.array([])
1✔
567
        self.starExtended = np.array([], dtype=int)
1✔
568
        self.lastDetected = np.empty((self.TargetList.nStars, 4), dtype=object)
1✔
569

570
    def __str__(self):
1✔
571
        """String representation of the Survey Simulation object
572

573
        When the command 'print' is used on the Survey Simulation object, this
574
        method will return the values contained in the object
575

576
        """
577

578
        for att in self.__dict__:
1✔
579
            print("%s: %r" % (att, getattr(self, att)))
1✔
580

581
        return "Survey Simulation class object attributes"
1✔
582

583
    def run_sim(self):
1✔
584
        """Performs the survey simulation"""
585

586
        OS = self.OpticalSystem
1✔
587
        TL = self.TargetList
1✔
588
        SU = self.SimulatedUniverse
1✔
589
        Obs = self.Observatory
1✔
590
        TK = self.TimeKeeping
1✔
591

592
        # TODO: start using this self.currentSep
593
        # set occulter separation if haveOcculter
594
        if OS.haveOcculter:
1✔
595
            self.currentSep = Obs.occulterSep
×
596

597
        # choose observing modes selected for detection (default marked with a flag)
598
        allModes = OS.observingModes
1✔
599
        det_mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
600
        # and for characterization (default is first spectro/IFS mode)
601
        spectroModes = list(
1✔
602
            filter(lambda mode: "spec" in mode["inst"]["name"], allModes)
603
        )
604
        if np.any(spectroModes):
1✔
605
            char_mode = spectroModes[0]
×
606
        # if no spectro mode, default char mode is first observing mode
607
        else:
608
            char_mode = allModes[0]
1✔
609

610
        # begin Survey, and loop until mission is finished
611
        log_begin = "OB%s: survey beginning." % (TK.OBnumber)
1✔
612
        self.logger.info(log_begin)
1✔
613
        self.vprint(log_begin)
1✔
614
        t0 = time.time()
1✔
615
        sInd = None
1✔
616
        ObsNum = 0
1✔
617
        while not TK.mission_is_over(OS, Obs, det_mode):
1✔
618
            # acquire the NEXT TARGET star index and create DRM
619
            old_sInd = sInd  # used to save sInd if returned sInd is None
1✔
620
            DRM, sInd, det_intTime, waitTime = self.next_target(sInd, det_mode)
1✔
621

622
            if sInd is not None:
1✔
623
                ObsNum += (
1✔
624
                    1  # we're making an observation so increment observation number
625
                )
626

627
                if OS.haveOcculter:
1✔
628
                    # advance to start of observation (add slew time for selected target
629
                    _ = TK.advanceToAbsTime(TK.currentTimeAbs.copy() + waitTime)
×
630

631
                # beginning of observation, start to populate DRM
632
                DRM["star_ind"] = sInd
1✔
633
                DRM["star_name"] = TL.Name[sInd]
1✔
634
                DRM["arrival_time"] = TK.currentTimeNorm.to("day").copy()
1✔
635
                DRM["OB_nb"] = TK.OBnumber
1✔
636
                DRM["ObsNum"] = ObsNum
1✔
637
                pInds = np.where(SU.plan2star == sInd)[0]
1✔
638
                DRM["plan_inds"] = pInds.astype(int)
1✔
639
                log_obs = (
1✔
640
                    "  Observation #%s, star ind %s (of %s) with %s planet(s), "
641
                    + "mission time at Obs start: %s, exoplanetObsTime: %s"
642
                ) % (
643
                    ObsNum,
644
                    sInd,
645
                    TL.nStars,
646
                    len(pInds),
647
                    TK.currentTimeNorm.to("day").copy().round(2),
648
                    TK.exoplanetObsTime.to("day").copy().round(2),
649
                )
650
                self.logger.info(log_obs)
1✔
651
                self.vprint(log_obs)
1✔
652

653
                # PERFORM DETECTION and populate revisit list attribute
654
                (
1✔
655
                    detected,
656
                    det_fZ,
657
                    det_JEZ,
658
                    det_systemParams,
659
                    det_SNR,
660
                    FA,
661
                ) = self.observation_detection(sInd, det_intTime.copy(), det_mode)
662
                # update the occulter wet mass
663
                if OS.haveOcculter:
1✔
664
                    DRM = self.update_occulter_mass(
×
665
                        DRM, sInd, det_intTime.copy(), "det"
666
                    )
667
                # populate the DRM with detection results
668
                DRM["det_time"] = det_intTime.to("day")
1✔
669
                DRM["det_status"] = detected
1✔
670
                DRM["det_SNR"] = det_SNR
1✔
671
                DRM["det_fZ"] = det_fZ.to(self.fZ_unit)
1✔
672
                DRM["det_JEZ"] = det_JEZ.to(self.JEZ_unit)
1✔
673
                DRM["det_params"] = det_systemParams
1✔
674

675
                # PERFORM CHARACTERIZATION and populate spectra list attribute
676
                if char_mode["SNR"] not in [0, np.inf]:
1✔
677
                    (
1✔
678
                        characterized,
679
                        char_fZ,
680
                        char_JEZ,
681
                        char_systemParams,
682
                        char_SNR,
683
                        char_intTime,
684
                    ) = self.observation_characterization(sInd, char_mode)
685
                else:
686
                    char_intTime = None
×
687
                    lenChar = len(pInds) + 1 if FA else len(pInds)
×
688
                    characterized = np.zeros(lenChar, dtype=float)
×
689
                    char_SNR = np.zeros(lenChar, dtype=float)
×
690
                    char_fZ = 0.0 / u.arcsec**2
×
691
                    char_systemParams = SU.dump_system_params(sInd)
×
692
                assert char_intTime != 0, "Integration time can't be 0."
1✔
693
                # update the occulter wet mass
694
                if OS.haveOcculter and (char_intTime is not None):
1✔
695
                    DRM = self.update_occulter_mass(DRM, sInd, char_intTime, "char")
×
696
                # populate the DRM with characterization results
697
                DRM["char_time"] = (
1✔
698
                    char_intTime.to("day") if char_intTime is not None else 0.0 * u.day
699
                )
700
                DRM["char_status"] = characterized[:-1] if FA else characterized
1✔
701
                DRM["char_SNR"] = char_SNR[:-1] if FA else char_SNR
1✔
702
                DRM["char_fZ"] = char_fZ.to(self.fZ_unit)
1✔
703
                DRM["char_params"] = char_systemParams
1✔
704
                # populate the DRM with FA results
705
                DRM["FA_det_status"] = int(FA)
1✔
706
                DRM["FA_char_status"] = characterized[-1] if FA else 0
1✔
707
                DRM["FA_char_SNR"] = char_SNR[-1] if FA else 0.0
1✔
708
                DRM["FA_char_JEZ"] = (
1✔
709
                    self.lastDetected[sInd, 1][-1] if FA else 0.0 * self.JEZ_unit
710
                )
711
                DRM["FA_char_dMag"] = self.lastDetected[sInd, 2][-1] if FA else 0.0
1✔
712
                DRM["FA_char_WA"] = (
1✔
713
                    self.lastDetected[sInd, 3][-1] * u.arcsec if FA else 0.0 * u.arcsec
714
                )
715

716
                # populate the DRM with observation modes
717
                DRM["det_mode"] = dict(det_mode)
1✔
718
                del DRM["det_mode"]["inst"], DRM["det_mode"]["syst"]
1✔
719
                DRM["char_mode"] = dict(char_mode)
1✔
720
                del DRM["char_mode"]["inst"], DRM["char_mode"]["syst"]
1✔
721
                DRM["exoplanetObsTime"] = TK.exoplanetObsTime.copy()
1✔
722

723
                # append result values to self.DRM
724
                self.DRM.append(DRM)
1✔
725

726
                # handle case of inf OBs and missionPortion < 1
727
                if np.isinf(TK.OBduration) and (TK.missionPortion < 1.0):
1✔
728
                    self.arbitrary_time_advancement(
1✔
729
                        TK.currentTimeNorm.to("day").copy() - DRM["arrival_time"]
730
                    )
731

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

836
    def arbitrary_time_advancement(self, dt):
1✔
837
        """Handles fully dynamically scheduled case where OBduration is infinite and
838
        missionPortion is less than 1.
839

840
        Args:
841
            dt (~astropy.units.quantity.Quantity):
842
                Total amount of time, including all overheads and extras used for the
843
                previous observation.
844

845
        Returns:
846
            None
847
        """
848

849
        self.TimeKeeping.allocate_time(
1✔
850
            dt
851
            * (1.0 - self.TimeKeeping.missionPortion)
852
            / self.TimeKeeping.missionPortion,
853
            addExoplanetObsTime=False,
854
        )
855

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

859
        This method chooses the next target star index based on which
860
        stars are available, their integration time, and maximum completeness.
861
        Returns None if no target could be found.
862

863
        Args:
864
            old_sInd (int):
865
                Index of the previous target star
866
            mode (dict):
867
                Selected observing mode for detection
868

869
        Returns:
870
            tuple:
871
                DRM (dict):
872
                    Design Reference Mission, contains the results of one complete
873
                    observation (detection and characterization)
874
                sInd (int):
875
                    Index of next target star. Defaults to None.
876
                intTime (astropy.units.Quantity):
877
                    Selected star integration time for detection in units of day.
878
                    Defaults to None.
879
                waitTime (astropy.units.Quantity):
880
                    a strategically advantageous amount of time to wait in the case
881
                    of an occulter for slew times
882

883
        """
884
        OS = self.OpticalSystem
1✔
885
        TL = self.TargetList
1✔
886
        Obs = self.Observatory
1✔
887
        TK = self.TimeKeeping
1✔
888

889
        # create DRM
890
        DRM = {}
1✔
891

892
        # allocate settling time + overhead time
893
        tmpCurrentTimeAbs = (
1✔
894
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
895
        )
896
        tmpCurrentTimeNorm = (
1✔
897
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
898
        )
899

900
        # create appropriate koMap
901
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
902

903
        # look for available targets
904
        # 1. initialize arrays
905
        slewTimes = np.zeros(TL.nStars) << u.d
1✔
906
        # fZs = np.zeros(TL.nStars) / u.arcsec**2.0
907
        dV = np.zeros(TL.nStars) << self.m_per_s
1✔
908
        intTimes = np.zeros(TL.nStars) << u.d
1✔
909
        obsTimes = np.zeros([2, TL.nStars]) << u.d
1✔
910
        sInds = np.arange(TL.nStars)
1✔
911

912
        # 2. find spacecraft orbital START positions (if occulter, positions
913
        # differ for each star) and filter out unavailable targets
914
        sd = None
1✔
915
        if OS.haveOcculter:
1✔
916
            sd = Obs.star_angularSep(TL, old_sInd, sInds, tmpCurrentTimeAbs)
×
917
            obsTimes = Obs.calculate_observableTimes(
×
918
                TL, sInds, tmpCurrentTimeAbs, self.koMaps, self.koTimes, mode
919
            )
920
            slewTimes = Obs.calculate_slewTimes(
×
921
                TL, old_sInd, sInds, sd, obsTimes, tmpCurrentTimeAbs
922
            )
923

924
        # 2.1 filter out totTimes > integration cutoff
925
        if len(sInds.tolist()) > 0:
1✔
926
            sInds = np.intersect1d(self.intTimeFilterInds, sInds)
1✔
927

928
        # start times, including slew times
929
        startTimes = tmpCurrentTimeAbs.copy() + slewTimes
1✔
930
        startTimesNorm = tmpCurrentTimeNorm.copy() + slewTimes
1✔
931

932
        # 2.5 Filter stars not observable at startTimes
933
        try:
1✔
934
            tmpIndsbool = list()
1✔
935
            for i in np.arange(len(sInds)):
1✔
936
                koTimeInd = np.where(
1✔
937
                    np.round(startTimes[sInds[i]].value) - self.koTimes.value == 0
938
                )[0][
939
                    0
940
                ]  # find indice where koTime is startTime[0]
941
                tmpIndsbool.append(
1✔
942
                    koMap[sInds[i]][koTimeInd].astype(bool)
943
                )  # Is star observable at time ind
944
            sInds = sInds[tmpIndsbool]
1✔
945
            del tmpIndsbool
1✔
946
        except:  # noqa: E722 # If there are no target stars to observe
×
947
            sInds = np.asarray([], dtype=int)
×
948

949
        # 3. filter out all previously (more-)visited targets, unless in
950
        if len(sInds.tolist()) > 0:
1✔
951
            sInds = self.revisitFilter(sInds, tmpCurrentTimeNorm)
1✔
952

953
        # 4.1 calculate integration times for ALL preselected targets
954
        (
1✔
955
            maxIntTimeOBendTime,
956
            maxIntTimeExoplanetObsTime,
957
            maxIntTimeMissionLife,
958
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
959
        maxIntTime = min(
1✔
960
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
961
        )  # Maximum intTime allowed
962

963
        if len(sInds.tolist()) > 0:
1✔
964
            if OS.haveOcculter and (old_sInd is not None):
1✔
965
                (
×
966
                    sInds,
967
                    slewTimes[sInds],
968
                    intTimes[sInds],
969
                    dV[sInds],
970
                ) = self.refineOcculterSlews(
971
                    old_sInd, sInds, slewTimes, obsTimes, sd, mode
972
                )
973
                endTimes = tmpCurrentTimeAbs.copy() + intTimes + slewTimes
×
974
            else:
975
                intTimes[sInds] = self.calc_targ_intTime(sInds, startTimes[sInds], mode)
1✔
976
                sInds = sInds[
1✔
977
                    np.where(intTimes[sInds] <= maxIntTime)
978
                ]  # Filters targets exceeding end of OB
979
                endTimes = tmpCurrentTimeAbs.copy() + intTimes
1✔
980

981
                if maxIntTime.value <= 0:
1✔
982
                    sInds = np.asarray([], dtype=int)
×
983

984
        # 5.1 TODO Add filter to filter out stars entering and exiting keepout
985
        # between startTimes and endTimes
986

987
        # 5.2 find spacecraft orbital END positions (for each candidate target),
988
        # and filter out unavailable targets
989
        if len(sInds.tolist()) > 0 and Obs.checkKeepoutEnd:
1✔
990
            # endTimes may exist past koTimes so we have an exception to hand this case
991
            try:
1✔
992
                tmpIndsbool = list()
1✔
993
                for i in np.arange(len(sInds)):
1✔
994
                    koTimeInd = np.where(
1✔
995
                        np.round(endTimes[sInds[i]].value) - self.koTimes.value == 0
996
                    )[0][
997
                        0
998
                    ]  # find indice where koTime is endTime[0]
999
                    tmpIndsbool.append(
1✔
1000
                        koMap[sInds[i]][koTimeInd].astype(bool)
1001
                    )  # Is star observable at time ind
1002
                sInds = sInds[tmpIndsbool]
1✔
1003
                del tmpIndsbool
1✔
1004
            except:  # noqa: E722
×
1005
                sInds = np.asarray([], dtype=int)
×
1006

1007
        # 6. choose best target from remaining
1008
        if len(sInds.tolist()) > 0:
1✔
1009
            # choose sInd of next target
1010
            sInd, waitTime = self.choose_next_target(
1✔
1011
                old_sInd, sInds, slewTimes, intTimes[sInds]
1012
            )
1013

1014
            # Should Choose Next Target decide there are no stars it wishes to
1015
            # observe at this time.
1016
            if sInd is None and (waitTime is not None):
1✔
1017
                self.vprint(
×
1018
                    "There are no stars available to observe. Waiting {}".format(
1019
                        waitTime
1020
                    )
1021
                )
1022
                return DRM, None, None, waitTime
×
1023
            elif (sInd is None) and (waitTime is None):
1✔
1024
                self.vprint(
×
1025
                    "There are no stars available to observe and waitTime is None."
1026
                )
1027
                return DRM, None, None, waitTime
×
1028
            # store selected star integration time
1029
            intTime = intTimes[sInd]
1✔
1030

1031
        # if no observable target, advanceTime to next Observable Target
1032
        else:
1033
            self.vprint(
×
1034
                "No Observable Targets at currentTimeNorm= "
1035
                + str(TK.currentTimeNorm.copy())
1036
            )
1037
            return DRM, None, None, None
×
1038

1039
        # update visited list for selected star
1040
        self.starVisits[sInd] += 1
1✔
1041
        # store normalized start time for future completeness update
1042
        self.lastObsTimes[sInd] = startTimesNorm[sInd]
1✔
1043

1044
        # populate DRM with occulter related values
1045
        if OS.haveOcculter:
1✔
1046
            DRM = Obs.log_occulterResults(
×
1047
                DRM, slewTimes[sInd], sInd, sd[sInd], dV[sInd]
1048
            )
1049
            return DRM, sInd, intTime, slewTimes[sInd]
×
1050

1051
        return DRM, sInd, intTime, waitTime
1✔
1052

1053
    def calc_targ_intTime(self, sInds, startTimes, mode):
1✔
1054
        """Helper method for next_target to aid in overloading for alternative
1055
        implementations.
1056

1057
        Given a subset of targets, calculate their integration times given the
1058
        start of observation time.
1059

1060
        Prototype just calculates integration times for fixed contrast depth.
1061

1062
        Args:
1063
            sInds (~numpy.ndarray(int)):
1064
                Indices of available targets
1065
            startTimes (astropy quantity array):
1066
                absolute start times of observations.
1067
                must be of the same size as sInds
1068
            mode (dict):
1069
                Selected observing mode for detection
1070

1071
        Returns:
1072
            ~astropy.units.Quantity(~numpy.ndarray(float)):
1073
                Integration times for detection
1074
                same dimension as sInds
1075

1076
        .. note::
1077

1078
            next_target filter will discard targets with zero integration times.
1079

1080
        """
1081

1082
        TL = self.TargetList
1✔
1083

1084
        # assumed values for detection
1085
        fZ = self.ZodiacalLight.fZ(
1✔
1086
            self.Observatory, self.TargetList, sInds, startTimes, mode
1087
        )
1088
        # Use the default exozodi intensities
1089
        JEZ = TL.JEZ0[mode["hex"]][sInds]
1✔
1090
        dMag = TL.int_dMag[sInds]
1✔
1091
        WA = TL.int_WA[sInds]
1✔
1092

1093
        # save out file containing photon count info
1094
        if self.record_counts_path is not None and len(self.count_lines) == 0:
1✔
1095
            C_p, C_b, C_sp, C_extra = self.OpticalSystem.Cp_Cb_Csp(
×
1096
                self.TargetList, sInds, fZ, JEZ, dMag, WA, mode, returnExtra=True
1097
            )
1098
            import csv
×
1099

1100
            count_fpath = os.path.join(self.record_counts_path, "counts")
×
1101

1102
            if not os.path.exists(count_fpath):
×
1103
                os.mkdir(count_fpath)
×
1104

1105
            outfile = os.path.join(count_fpath, str(self.seed) + ".csv")
×
1106
            self.count_lines.append(
×
1107
                [
1108
                    "sInd",
1109
                    "HIPs",
1110
                    "C_F0",
1111
                    "C_p0",
1112
                    "C_sr",
1113
                    "C_z",
1114
                    "C_ez",
1115
                    "C_dc",
1116
                    "C_cc",
1117
                    "C_rn",
1118
                    "C_p",
1119
                    "C_b",
1120
                    "C_sp",
1121
                ]
1122
            )
1123

1124
            for i, sInd in enumerate(sInds):
×
1125
                self.count_lines.append(
×
1126
                    [
1127
                        sInd,
1128
                        self.TargetList.Name[sInd],
1129
                        C_extra["C_F0"][0].value,
1130
                        C_extra["C_sr"][i].value,
1131
                        C_extra["C_z"][i].value,
1132
                        C_extra["C_ez"][i].value,
1133
                        C_extra["C_dc"][i].value,
1134
                        C_extra["C_cc"][i].value,
1135
                        C_extra["C_rn"][i].value,
1136
                        C_p[i].value,
1137
                        C_b[i].value,
1138
                        C_sp[i].value,
1139
                    ]
1140
                )
1141

1142
            with open(outfile, "w") as csvfile:
×
1143
                c = csv.writer(csvfile)
×
1144
                c.writerows(self.count_lines)
×
1145

1146
        intTimes = self.OpticalSystem.calc_intTime(
1✔
1147
            self.TargetList, sInds, fZ, JEZ, dMag, WA, mode
1148
        )
1149
        intTimes[~np.isfinite(intTimes)] = 0 * u.d
1✔
1150

1151
        return intTimes
1✔
1152

1153
    def choose_next_target(self, old_sInd, sInds, slewTimes, intTimes):
1✔
1154
        """Helper method for method next_target to simplify alternative implementations.
1155

1156
        Given a subset of targets (pre-filtered by method next_target or some
1157
        other means), select the best next one. The prototype uses completeness
1158
        as the sole heuristic.
1159

1160
        Args:
1161
            old_sInd (int):
1162
                Index of the previous target star
1163
            sInds (~numpy.ndarray(int)):
1164
                Indices of available targets
1165
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1166
                slew times to all stars (must be indexed by sInds)
1167
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1168
                Integration times for detection in units of day
1169

1170
        Returns:
1171
            tuple:
1172
                sInd (int):
1173
                    Index of next target star
1174
                waitTime (:py:class:`~astropy.units.Quantity`):
1175
                    Some strategic amount of time to wait in case an occulter slew is
1176
                    desired (default is None)
1177

1178
        """
1179

1180
        Comp = self.Completeness
1✔
1181
        TL = self.TargetList
1✔
1182
        TK = self.TimeKeeping
1✔
1183
        OS = self.OpticalSystem
1✔
1184
        Obs = self.Observatory
1✔
1185
        allModes = OS.observingModes
1✔
1186

1187
        # cast sInds to array
1188
        sInds = np.array(sInds, ndmin=1, copy=copy_if_needed)
1✔
1189
        # calculate dt since previous observation
1190
        dt = TK.currentTimeNorm.copy() + slewTimes[sInds] - self.lastObsTimes[sInds]
1✔
1191
        # get dynamic completeness values
1192
        comps = Comp.completeness_update(TL, sInds, self.starVisits[sInds], dt)
1✔
1193
        # choose target with maximum completeness
1194
        sInd = np.random.choice(sInds[comps == max(comps)])
1✔
1195

1196
        # Check if exoplanetObsTime would be exceeded
1197
        mode = list(filter(lambda mode: mode["detectionMode"], allModes))[0]
1✔
1198
        (
1✔
1199
            maxIntTimeOBendTime,
1200
            maxIntTimeExoplanetObsTime,
1201
            maxIntTimeMissionLife,
1202
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode)
1203
        maxIntTime = min(
1✔
1204
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1205
        )  # Maximum intTime allowed
1206
        intTimes2 = self.calc_targ_intTime(
1✔
1207
            np.array([sInd]), TK.currentTimeAbs.copy(), mode
1208
        )
1209
        if (
1✔
1210
            intTimes2 > maxIntTime
1211
        ):  # check if max allowed integration time would be exceeded
1212
            self.vprint("max allowed integration time would be exceeded")
×
1213
            sInd = None
×
1214
            waitTime = 1.0 * u.d
×
1215
            return sInd, waitTime
×
1216

1217
        return (
1✔
1218
            sInd,
1219
            slewTimes[sInd],
1220
        )  # if coronagraph or first sInd, waitTime will be 0 days
1221

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

1225
        Refines the selection of occulter slew times by filtering based on mission time
1226
        constraints and selecting the best slew time for each star. This method calls on
1227
        other occulter methods within SurveySimulation depending on how slew times were
1228
        calculated prior to calling this function (i.e. depending on which
1229
        implementation of the Observatory module is used).
1230

1231
        Args:
1232
            old_sInd (int):
1233
                Index of the previous target star
1234
            sInds (~numpy.ndarray(int)):
1235
                Indices of available targets
1236
            slewTimes (astropy quantity array):
1237
                slew times to all stars (must be indexed by sInds)
1238
            obsTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1239
                A binary array with TargetList.nStars rows and
1240
                (missionFinishAbs-missionStart)/dt columns
1241
                where dt is 1 day by default. A value of 1 indicates the star is in
1242
                keepout (and therefore cannot be observed). A value of 0 indicates the
1243
                star is not in keepout and may be observed.
1244
            sd (~astropy.units.Quantity(~numpy.ndarray(float))):
1245
                Angular separation between stars in rad
1246
            mode (dict):
1247
                Selected observing mode for detection
1248

1249
        Returns:
1250
            tuple:
1251
                sInds (int):
1252
                    Indeces of next target star
1253
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1254
                    slew times to all stars (must be indexed by sInds)
1255
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1256
                    Integration times for detection in units of day
1257
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1258
                    Delta-V used to transfer to new star line of sight in unis of m/s
1259
        """
1260

1261
        Obs = self.Observatory
×
1262
        TL = self.TargetList
×
1263

1264
        # initializing arrays
1265
        obsTimeArray = np.zeros([TL.nStars, 50]) << u.d
×
1266
        intTimeArray = np.zeros([TL.nStars, 2]) << u.d
×
1267

1268
        for n in sInds:
×
1269
            obsTimeArray[n, :] = (
×
1270
                np.linspace(obsTimes[0, n].value, obsTimes[1, n].value, 50) << u.d
1271
            )
1272
        intTimeArray[sInds, 0] = self.calc_targ_intTime(
×
1273
            sInds, Time(obsTimeArray[sInds, 0], format="mjd", scale="tai"), mode
1274
        )
1275
        intTimeArray[sInds, 1] = self.calc_targ_intTime(
×
1276
            sInds, Time(obsTimeArray[sInds, -1], format="mjd", scale="tai"), mode
1277
        )
1278

1279
        # determining which scheme to use to filter slews
1280
        obsModName = Obs.__class__.__name__
×
1281

1282
        # slew times have not been calculated/decided yet (SotoStarshade)
1283
        if obsModName == "SotoStarshade":
×
1284
            sInds, intTimes, slewTimes, dV = self.findAllowableOcculterSlews(
×
1285
                sInds,
1286
                old_sInd,
1287
                sd[sInds],
1288
                slewTimes[sInds],
1289
                obsTimeArray[sInds, :],
1290
                intTimeArray[sInds, :],
1291
                mode,
1292
            )
1293

1294
        # slew times were calculated/decided beforehand (Observatory Prototype)
1295
        else:
1296
            sInds, intTimes, slewTimes = self.filterOcculterSlews(
×
1297
                sInds,
1298
                slewTimes[sInds],
1299
                obsTimeArray[sInds, :],
1300
                intTimeArray[sInds, :],
1301
                mode,
1302
            )
1303
            dV = np.zeros(len(sInds)) << self.m_per_s
×
1304

1305
        return sInds, slewTimes, intTimes, dV
×
1306

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

1310
        Used by the refineOcculterSlews method when slew times have been selected
1311
        a priori. This method filters out slews that are not within desired observing
1312
        blocks, the maximum allowed integration time, and are outside of future
1313
        keepouts.
1314

1315
        Args:
1316
            sInds (~numpy.ndarray(int)):
1317
                Indices of available targets
1318
            slewTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1319
                slew times to all stars (must be indexed by sInds)
1320
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1321
                Array of times during which a star is out of keepout, has shape
1322
                nx50 where n is the number of stars in sInds. Unit of days
1323
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1324
                Array of integration times for each time in obsTimeArray, has shape
1325
                nx2 where n is the number of stars in sInds. Unit of days
1326
            mode (dict):
1327
                Selected observing mode for detection
1328

1329
        Returns:
1330
            tuple:
1331
                sInds (int):
1332
                    Indeces of next target star
1333
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1334
                    Integration times for detection in units of day
1335
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1336
                    slew times to all stars (must be indexed by sInds)
1337
        """
1338

1339
        TK = self.TimeKeeping
×
1340
        Obs = self.Observatory
×
1341

1342
        # allocate settling time + overhead time
1343
        tmpCurrentTimeAbs = (
×
1344
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1345
        )
1346
        tmpCurrentTimeNorm = (
×
1347
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1348
        )
1349

1350
        # 0. lambda function that linearly interpolates Integration Time
1351
        # between obsTimes
1352
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1353
            t - np.array(x[:, 0]).reshape(len(t), 1)
1354
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1355

1356
        # 1. initializing arrays
1357
        obsTimesRange = np.array(
×
1358
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1359
        )  # nx2 array with start and end times of obsTimes for each star
1360
        intTimesRange = np.array([intTimeArray[:, 0], intTimeArray[:, -1]])
×
1361

1362
        OBnumbers = np.zeros(
×
1363
            [len(sInds), 1]
1364
        )  # for each sInd, will show during which OB observations will take place
1365
        maxIntTimes = np.zeros([len(sInds), 1]) << u.d
×
1366

1367
        intTimes = (
×
1368
            linearInterp(
1369
                intTimesRange.T,
1370
                obsTimesRange.T,
1371
                (tmpCurrentTimeAbs + slewTimes).reshape(len(sInds), 1).value,
1372
            )
1373
            << u.d
1374
        )  # calculate intTimes for each slew time
1375

1376
        minObsTimeNorm = (obsTimesRange[0, :] - tmpCurrentTimeAbs.value).reshape(
×
1377
            [len(sInds), 1]
1378
        )
1379
        maxObsTimeNorm = (obsTimesRange[1, :] - tmpCurrentTimeAbs.value).reshape(
×
1380
            [len(sInds), 1]
1381
        )
1382
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1383

1384
        # 2. find OBnumber for each sInd's slew time
1385
        if len(TK.OBendTimes) > 1:
×
1386
            for i in range(len(sInds)):
×
1387
                S = np.where(
×
1388
                    TK.OBstartTimes.value - tmpCurrentTimeNorm.value
1389
                    < slewTimes[i].value
1390
                )[0][-1]
1391
                F = np.where(
×
1392
                    TK.OBendTimes.value - tmpCurrentTimeNorm.value < slewTimes[i].value
1393
                )[0]
1394

1395
                # case when slews are in the first OB
1396
                if F.shape[0] == 0:
×
1397
                    F = -1
×
1398
                else:
1399
                    F = F[-1]
×
1400

1401
                # slew occurs within an OB (nth OB has started but hasn't ended)
1402
                if S != F:
×
1403
                    OBnumbers[i] = S
×
1404
                    (
×
1405
                        maxIntTimeOBendTime,
1406
                        maxIntTimeExoplanetObsTime,
1407
                        maxIntTimeMissionLife,
1408
                    ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, TK.OBstartTimes[S], S)
1409
                    maxIntTimes[i] = min(
×
1410
                        maxIntTimeOBendTime,
1411
                        maxIntTimeExoplanetObsTime,
1412
                        maxIntTimeMissionLife,
1413
                    )  # Maximum intTime allowed
1414

1415
                # slew occurs between OBs, badbadnotgood
1416
                else:
1417
                    OBnumbers[i] = -1
×
1418
                    maxIntTimes[i] = 0 * u.d
×
1419
            OBstartTimeNorm = (
×
1420
                TK.OBstartTimes[np.array(OBnumbers, dtype=int)].value
1421
                - tmpCurrentTimeNorm.value
1422
            )
1423
        else:
1424
            (
×
1425
                maxIntTimeOBendTime,
1426
                maxIntTimeExoplanetObsTime,
1427
                maxIntTimeMissionLife,
1428
            ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm)
1429
            maxIntTimes[:] = min(
×
1430
                maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1431
            )  # Maximum intTime allowed
1432
            OBstartTimeNorm = np.zeros(OBnumbers.shape)
×
1433

1434
        # finding the minimum possible slew time
1435
        # (either OBstart or when star JUST leaves keepout)
1436
        minAllowedSlewTime = np.max([OBstartTimeNorm, minObsTimeNorm], axis=0)
×
1437

1438
        # 3. start filtering process
1439
        good_inds = np.where((OBnumbers >= 0) & (ObsTimeRange > intTimes.value))[0]
×
1440
        # star slew within OB -AND- can finish observing
1441
        # before it goes back into keepout
1442
        if good_inds.shape[0] > 0:
×
1443
            # the good ones
1444
            sInds = sInds[good_inds]
×
1445
            slewTimes = slewTimes[good_inds]
×
1446
            intTimes = intTimes[good_inds]
×
1447
            OBstartTimeNorm = OBstartTimeNorm[good_inds]
×
1448
            minAllowedSlewTime = minAllowedSlewTime[good_inds]
×
1449

1450
            # maximum allowed slew time based on integration times
1451
            maxAllowedSlewTime = maxIntTimes[good_inds].value - intTimes.value
×
1452
            maxAllowedSlewTime[maxAllowedSlewTime < 0] = -np.inf
×
1453
            maxAllowedSlewTime += OBstartTimeNorm  # calculated rel to currentTime norm
×
1454

1455
            # checking to see if slewTimes are allowed
1456
            good_inds = np.where(
×
1457
                (slewTimes.reshape([len(sInds), 1]).value > minAllowedSlewTime)
1458
                & (slewTimes.reshape([len(sInds), 1]).value < maxAllowedSlewTime)
1459
            )[0]
1460

1461
            slewTimes = slewTimes[good_inds]
×
1462
        else:
1463
            slewTimes = slewTimes[good_inds]
×
1464

1465
        return sInds[good_inds], intTimes[good_inds].flatten(), slewTimes
×
1466

1467
    def findAllowableOcculterSlews(
1✔
1468
        self, sInds, old_sInd, sd, slewTimes, obsTimeArray, intTimeArray, mode
1469
    ):
1470
        """Finds an array of allowable slew times for each star
1471

1472
        Used by the refineOcculterSlews method when slew times have NOT been selected
1473
        a priori. This method creates nx50 arrays (where the row corresponds to a
1474
        specific star and the column corresponds to a future point in time relative to
1475
        currentTime).
1476

1477
        These arrays are initially zero but are populated with the corresponding values
1478
        (slews, intTimes, etc) if slewing to that time point (i.e. beginning an
1479
        observation) would lead to a successful observation. A "successful observation"
1480
        is defined by certain conditions relating to keepout and the komap, observing
1481
        blocks, mission lifetime, and some constraints on the dVmap calculation in
1482
        SotoStarshade. Each star will likely have a range of slewTimes that would lead
1483
        to a successful observation -- another method is then called to select the best
1484
        of these slewTimes.
1485

1486
        Args:
1487
            sInds (~numpy.ndarray(int)):
1488
                Indices of available targets
1489
            old_sInd (int):
1490
                Index of the previous target star
1491
            sd (~astropy.units.Quantity(~numpy.ndarray(float))):
1492
                Angular separation between stars in rad
1493
            slewTimes (astropy quantity array):
1494
                slew times to all stars (must be indexed by sInds)
1495
            obsTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1496
                Array of times during which a star is out of keepout, has shape
1497
                nx50 where n is the number of stars in sInds
1498
            intTimeArray (~astropy.units.Quantity(~numpy.ndarray(float))):
1499
                Array of integration times for each time in obsTimeArray, has shape
1500
                nx50 where n is the number of stars in sInds
1501
            mode (dict):
1502
                Selected observing mode for detection
1503

1504
        Returns:
1505
            tuple:
1506
                sInds (numpy.ndarray(int)):
1507
                    Indices of next target star
1508
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1509
                    slew times to all stars (must be indexed by sInds)
1510
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1511
                    Integration times for detection in units of day
1512
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1513
                    Delta-V used to transfer to new star line of sight in unis of m/s
1514
        """
1515
        TK = self.TimeKeeping
×
1516
        Obs = self.Observatory
×
1517
        TL = self.TargetList
×
1518

1519
        # 0. lambda function that linearly interpolates Integration
1520
        # Time between obsTimes
1521
        linearInterp = lambda y, x, t: np.diff(y) / np.diff(x) * (  # noqa: E731
×
1522
            t - np.array(x[:, 0]).reshape(len(t), 1)
1523
        ) + np.array(y[:, 0]).reshape(len(t), 1)
1524

1525
        # allocate settling time + overhead time
1526
        tmpCurrentTimeAbs = (
×
1527
            TK.currentTimeAbs.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1528
        )
1529
        tmpCurrentTimeNorm = (
×
1530
            TK.currentTimeNorm.copy() + Obs.settlingTime + mode["syst"]["ohTime"]
1531
        )
1532

1533
        # 1. initializing arrays
1534
        obsTimes = np.array(
×
1535
            [obsTimeArray[:, 0], obsTimeArray[:, -1]]
1536
        )  # nx2 array with start and end times of obsTimes for each star
1537
        intTimes_int = (
×
1538
            np.zeros(obsTimeArray.shape) << u.d
1539
        )  # initializing intTimes of shape nx50 then interpolating
1540
        intTimes_int = (
×
1541
            np.hstack(
1542
                [
1543
                    intTimeArray[:, 0].reshape(len(sInds), 1).value,
1544
                    linearInterp(
1545
                        intTimeArray.value, obsTimes.T, obsTimeArray[:, 1:].value
1546
                    ),
1547
                ]
1548
            )
1549
            << u.d
1550
        )
1551
        allowedSlewTimes = np.zeros(obsTimeArray.shape) << u.d
×
1552
        allowedintTimes = np.zeros(obsTimeArray.shape) << u.d
×
1553
        allowedCharTimes = np.zeros(obsTimeArray.shape) << u.d
×
1554
        obsTimeArrayNorm = obsTimeArray.value - tmpCurrentTimeAbs.value
×
1555

1556
        # obsTimes -> relative to current Time
1557
        try:
×
1558
            minObsTimeNorm = np.array([np.min(v[v > 0]) for v in obsTimeArrayNorm])
×
1559
        except:  # noqa: E722
×
1560
            # an error pops up sometimes at the end of the mission, this fixes it
1561
            # TODO: define the error type that occurs
1562
            # rewrite to avoid a try/except if possible
1563
            minObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1564

1565
        maxObsTimeNorm = obsTimes[1, :].T - tmpCurrentTimeAbs.value
×
1566
        ObsTimeRange = maxObsTimeNorm - minObsTimeNorm
×
1567

1568
        # getting max possible intTime
1569
        (
×
1570
            maxIntTimeOBendTime,
1571
            maxIntTimeExoplanetObsTime,
1572
            maxIntTimeMissionLife,
1573
        ) = TK.get_ObsDetectionMaxIntTime(Obs, mode, tmpCurrentTimeNorm, TK.OBnumber)
1574
        maxIntTime = min(
×
1575
            maxIntTimeOBendTime, maxIntTimeExoplanetObsTime, maxIntTimeMissionLife
1576
        )  # Maximum intTime allowed
1577

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

1582
        # first filled in for the current OB
1583
        minAllowedSlewTimes = np.array(
×
1584
            [minObsTimeNorm.T] * len(intTimes_int.T)
1585
        ).T  # just to make it nx50 so it plays nice with the other arrays
1586
        maxAllowedSlewTimes = maxIntTime.value - intTimes_int.value
×
1587
        maxAllowedSlewTimes[maxAllowedSlewTimes > Obs.occ_dtmax.value] = (
×
1588
            Obs.occ_dtmax.value
1589
        )
1590

1591
        # conditions that must be met to define an allowable slew time
1592
        cond1 = (
×
1593
            minAllowedSlewTimes >= Obs.occ_dtmin.value
1594
        )  # minimum dt time in dV map interpolant
1595
        cond2 = (
×
1596
            maxAllowedSlewTimes <= Obs.occ_dtmax.value
1597
        )  # maximum dt time in dV map interpolant
1598
        cond3 = maxAllowedSlewTimes > minAllowedSlewTimes
×
1599
        cond4 = intTimes_int.value < ObsTimeRange.reshape(len(sInds), 1)
×
1600

1601
        conds = cond1 & cond2 & cond3 & cond4
×
1602
        minAllowedSlewTimes[np.invert(conds)] = (
×
1603
            np.inf
1604
        )  # these are filtered during the next filter
1605
        maxAllowedSlewTimes[np.invert(conds)] = -np.inf
×
1606

1607
        # one last condition to meet
1608
        map_i, map_j = np.where(
×
1609
            (obsTimeArrayNorm > minAllowedSlewTimes)
1610
            & (obsTimeArrayNorm < maxAllowedSlewTimes)
1611
        )
1612

1613
        # 2.5 if any stars are slew-able to within this OB block, populate
1614
        # "allowedSlewTimes", a running tally of possible slews
1615
        # within the time range a star is observable (out of keepout)
1616
        if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1617
            allowedSlewTimes[map_i, map_j] = obsTimeArrayNorm[map_i, map_j] * u.d
×
1618
            allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1619
            allowedCharTimes[map_i, map_j] = maxIntTime - intTimes_int[map_i, map_j]
×
1620

1621
        # 3. search future OBs
1622
        OB_withObsStars = (
×
1623
            TK.OBstartTimes.value - np.min(obsTimeArrayNorm) - tmpCurrentTimeNorm.value
1624
        )  # OBs within which any star is observable
1625

1626
        if any(OB_withObsStars > 0):
×
1627
            nOBstart = np.argmin(np.abs(OB_withObsStars))
×
1628
            nOBend = np.argmax(OB_withObsStars)
×
1629

1630
            # loop through the next 5 OBs (or until mission is over if there are less
1631
            # than 5 OBs in the future)
1632
            for i in np.arange(nOBstart, np.min([nOBend, nOBstart + 5])):
×
1633
                # max int Times for the next OB
1634
                (
×
1635
                    maxIntTimeOBendTime,
1636
                    maxIntTimeExoplanetObsTime,
1637
                    maxIntTimeMissionLife,
1638
                ) = TK.get_ObsDetectionMaxIntTime(
1639
                    Obs, mode, TK.OBstartTimes[i + 1], i + 1
1640
                )
1641
                maxIntTime_nOB = min(
×
1642
                    maxIntTimeOBendTime,
1643
                    maxIntTimeExoplanetObsTime,
1644
                    maxIntTimeMissionLife,
1645
                )  # Maximum intTime allowed
1646

1647
                # min and max slew times rel to current Time (norm)
1648
                nOBstartTimeNorm = np.array(
×
1649
                    [TK.OBstartTimes[i + 1].value - tmpCurrentTimeNorm.value]
1650
                    * len(sInds)
1651
                )
1652

1653
                # min slew times for stars start either whenever the star first leaves
1654
                # keepout or when next OB stars, whichever comes last
1655
                minAllowedSlewTimes_nOB = np.array(
×
1656
                    [np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T]
1657
                    * len(maxAllowedSlewTimes.T)
1658
                ).T
1659
                maxAllowedSlewTimes_nOB = (
×
1660
                    nOBstartTimeNorm.reshape(len(sInds), 1)
1661
                    + maxIntTime_nOB.value
1662
                    - intTimes_int.value
1663
                )
1664
                maxAllowedSlewTimes_nOB[
×
1665
                    maxAllowedSlewTimes_nOB > Obs.occ_dtmax.value
1666
                ] = Obs.occ_dtmax.value
1667

1668
                # amount of time left when the stars are still out of keepout
1669
                ObsTimeRange_nOB = (
×
1670
                    maxObsTimeNorm
1671
                    - np.max([minObsTimeNorm, nOBstartTimeNorm], axis=0).T
1672
                )
1673

1674
                # condition to be met for an allowable slew time
1675
                cond1 = minAllowedSlewTimes_nOB >= Obs.occ_dtmin.value
×
1676
                cond2 = maxAllowedSlewTimes_nOB <= Obs.occ_dtmax.value
×
1677
                cond3 = maxAllowedSlewTimes_nOB > minAllowedSlewTimes_nOB
×
1678
                cond4 = intTimes_int.value < ObsTimeRange_nOB.reshape(len(sInds), 1)
×
1679
                cond5 = intTimes_int.value < maxIntTime_nOB.value
×
1680
                conds = cond1 & cond2 & cond3 & cond4 & cond5
×
1681

1682
                minAllowedSlewTimes_nOB[np.invert(conds)] = np.inf
×
1683
                maxAllowedSlewTimes_nOB[np.invert(conds)] = -np.inf
×
1684

1685
                # one last condition
1686
                map_i, map_j = np.where(
×
1687
                    (obsTimeArrayNorm > minAllowedSlewTimes_nOB)
1688
                    & (obsTimeArrayNorm < maxAllowedSlewTimes_nOB)
1689
                )
1690

1691
                # 3.33 populate the running tally of allowable slew times if it meets
1692
                # all conditions
1693
                if map_i.shape[0] > 0 and map_j.shape[0] > 0:
×
1694
                    allowedSlewTimes[map_i, map_j] = (
×
1695
                        obsTimeArrayNorm[map_i, map_j] * u.d
1696
                    )
1697
                    allowedintTimes[map_i, map_j] = intTimes_int[map_i, map_j]
×
1698
                    allowedCharTimes[map_i, map_j] = (
×
1699
                        maxIntTime_nOB - intTimes_int[map_i, map_j]
1700
                    )
1701

1702
        # 3.67 filter out any stars that are not observable at all
1703
        filterDuds = np.sum(allowedSlewTimes, axis=1) > 0.0
×
1704
        sInds = sInds[filterDuds]
×
1705

1706
        # 4. choose a slew time for each available star
1707
        # calculate dVs for each possible slew time for each star
1708
        allowed_dVs = Obs.calculate_dV(
×
1709
            TL,
1710
            old_sInd,
1711
            sInds,
1712
            sd[filterDuds],
1713
            allowedSlewTimes[filterDuds, :],
1714
            tmpCurrentTimeAbs,
1715
        )
1716

1717
        if len(sInds.tolist()) > 0:
×
1718
            # select slew time for each star
1719
            dV_inds = np.arange(0, len(sInds))
×
1720
            sInds, intTime, slewTime, dV = self.chooseOcculterSlewTimes(
×
1721
                sInds,
1722
                allowedSlewTimes[filterDuds, :],
1723
                allowed_dVs[dV_inds, :],
1724
                allowedintTimes[filterDuds, :],
1725
                allowedCharTimes[filterDuds, :],
1726
            )
1727

1728
            return sInds, intTime, slewTime, dV
×
1729

1730
        else:
1731
            empty = np.asarray([], dtype=int)
×
1732
            return empty, empty << u.d, empty << u.d, empty << self.m_per_s
×
1733

1734
    def chooseOcculterSlewTimes(self, sInds, slewTimes, dV, intTimes, charTimes):
1✔
1735
        """Selects the best slew time for each star
1736

1737
        This method searches through an array of permissible slew times for
1738
        each star and chooses the best slew time for the occulter based on
1739
        maximizing possible characterization time for that particular star (as
1740
        a default).
1741

1742
        Args:
1743
            sInds (~numpy.ndarray(int)):
1744
                Indices of available targets
1745
            slewTimes (astropy quantity array):
1746
                slew times to all stars (must be indexed by sInds)
1747
            dV (~astropy.units.Quantity(~numpy.ndarray(float))):
1748
                Delta-V used to transfer to new star line of sight in unis of m/s
1749
            intTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1750
                Integration times for detection in units of day
1751
            charTimes (~astropy.units.Quantity(~numpy.ndarray(float))):
1752
                Time left over after integration which could be used for
1753
                characterization in units of day
1754

1755
        Returns:
1756
            tuple:
1757
                sInds (int):
1758
                    Indeces of next target star
1759
                slewTimes (astropy.units.Quantity(numpy.ndarray(float))):
1760
                    slew times to all stars (must be indexed by sInds)
1761
                intTimes (astropy.units.Quantity(numpy.ndarray(float))):
1762
                    Integration times for detection in units of day
1763
                dV (astropy.units.Quantity(numpy.ndarray(float))):
1764
                    Delta-V used to transfer to new star line of sight in unis of m/s
1765
        """
1766

1767
        # selection criteria for each star slew
1768
        good_j = np.argmax(
×
1769
            charTimes, axis=1
1770
        )  # maximum possible characterization time available
1771
        good_i = np.arange(0, len(sInds))
×
1772

1773
        dV = dV[good_i, good_j]
×
1774
        intTime = intTimes[good_i, good_j]
×
1775
        slewTime = slewTimes[good_i, good_j]
×
1776

1777
        return sInds, intTime, slewTime, dV
×
1778

1779
    def observation_detection(self, sInd, intTime, mode):
1✔
1780
        """Determines SNR and detection status for a given integration time
1781
        for detection. Also updates the lastDetected and starRevisit lists.
1782

1783
        Args:
1784
            sInd (int):
1785
                Integer index of the star of interest
1786
            intTime (~astropy.units.Quantity(~numpy.ndarray(float))):
1787
                Selected star integration time for detection in units of day.
1788
                Defaults to None.
1789
            mode (dict):
1790
                Selected observing mode for detection
1791

1792
        Returns:
1793
            tuple:
1794
                detected (numpy.ndarray(int)):
1795
                    Detection status for each planet orbiting the observed target star:
1796
                    1 is detection, 0 missed detection, -1 below IWA, and -2 beyond OWA
1797
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
1798
                    Surface brightness of local zodiacal light in units of 1/arcsec2
1799
                JEZ (astropy.units.Quantity(numpy.ndarray(float))):
1800
                    Intensity of exo-zodiacal light in units of photons/s/m2/arcsec2
1801
                systemParams (dict):
1802
                    Dictionary of time-dependant planet properties averaged over the
1803
                    duration of the integration
1804
                SNR (numpy.darray(float)):
1805
                    Detection signal-to-noise ratio of the observable planets
1806
                FA (bool):
1807
                    False alarm (false positive) boolean
1808

1809
        """
1810

1811
        PPop = self.PlanetPopulation
1✔
1812
        ZL = self.ZodiacalLight
1✔
1813
        PPro = self.PostProcessing
1✔
1814
        TL = self.TargetList
1✔
1815
        SU = self.SimulatedUniverse
1✔
1816
        Obs = self.Observatory
1✔
1817
        TK = self.TimeKeeping
1✔
1818

1819
        # Save Current Time before attempting time allocation
1820
        currentTimeNorm = TK.currentTimeNorm.copy()
1✔
1821
        currentTimeAbs = TK.currentTimeAbs.copy()
1✔
1822

1823
        # Allocate Time
1824
        extraTime = intTime * (mode["timeMultiplier"] - 1.0)  # calculates extraTime
1✔
1825
        success = TK.allocate_time(
1✔
1826
            intTime + extraTime + Obs.settlingTime + mode["syst"]["ohTime"], True
1827
        )
1828
        assert success, "Could not allocate observation detection time ({}).".format(
1✔
1829
            intTime + extraTime + Obs.settlingTime + mode["syst"]["ohTime"]
1830
        )
1831
        # calculates partial time to be added for every ntFlux
1832
        dt = intTime / float(self.ntFlux)
1✔
1833
        # find indices of planets around the target
1834
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
1835

1836
        # initialize outputs
1837
        detected = np.array([], dtype=int)
1✔
1838
        # write current system params by default
1839
        systemParams = SU.dump_system_params(sInd)
1✔
1840
        SNR = np.zeros(len(pInds))
1✔
1841

1842
        # if any planet, calculate SNR
1843
        if len(pInds) > 0:
1✔
1844
            # initialize arrays for SNR integration
1845
            fZs = np.zeros(self.ntFlux) << self.fZ_unit
1✔
1846
            systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
1847
            JEZs = np.zeros((self.ntFlux, len(pInds))) << self.JEZ_unit
1✔
1848
            Ss = np.zeros((self.ntFlux, len(pInds)))
1✔
1849
            Ns = np.zeros((self.ntFlux, len(pInds)))
1✔
1850
            # accounts for the time since the current time
1851
            timePlus = Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
1✔
1852
            # integrate the signal (planet flux) and noise
1853
            for i in range(self.ntFlux):
1✔
1854
                # allocate first half of dt
1855
                timePlus += dt / 2.0
1✔
1856
                # calculate current zodiacal light brightness
1857
                fZs[i] = ZL.fZ(
1✔
1858
                    Obs,
1859
                    TL,
1860
                    np.array([sInd], ndmin=1),
1861
                    (currentTimeAbs + timePlus).reshape(1),
1862
                    mode,
1863
                )[0]
1864
                # propagate the system to match up with current time
1865
                SU.propag_system(
1✔
1866
                    sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
1867
                )
1868
                # Calculate the exozodi intensity
1869
                JEZs[i] = SU.scale_JEZ(sInd, mode)
1✔
1870
                self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
1871
                # save planet parameters
1872
                systemParamss[i] = SU.dump_system_params(sInd)
1✔
1873
                # calculate signal and noise (electron count rates)
1874
                Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
1875
                    sInd, pInds, dt, mode, fZ=fZs[i], JEZ=JEZs[i]
1876
                )
1877
                # allocate second half of dt
1878
                timePlus += dt / 2.0
1✔
1879

1880
            # average output parameters
1881
            fZ = np.mean(fZs)
1✔
1882
            JEZ = np.mean(JEZs, axis=0)
1✔
1883
            systemParams = {
1✔
1884
                key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
1885
                / float(self.ntFlux)
1886
                for key in sorted(systemParamss[0])
1887
            }
1888
            # calculate SNR
1889
            S = Ss.sum(0)
1✔
1890
            N = Ns.sum(0)
1✔
1891
            SNR[N > 0] = S[N > 0] / N[N > 0]
1✔
1892

1893
        # if no planet, just save zodiacal brightness in the middle of the integration
1894
        else:
1895
            totTime = intTime * (mode["timeMultiplier"])
1✔
1896
            fZ = ZL.fZ(
1✔
1897
                Obs,
1898
                TL,
1899
                np.array([sInd], ndmin=1),
1900
                (currentTimeAbs + totTime / 2.0).reshape(1),
1901
                mode,
1902
            )[0]
1903
            # Use the default star value if no planets
1904
            JEZ = TL.JEZ0[mode["hex"]][sInd]
1✔
1905

1906
        # find out if a false positive (false alarm) or any false negative
1907
        # (missed detections) have occurred
1908
        FA, MD = PPro.det_occur(SNR, mode, TL, sInd, intTime)
1✔
1909

1910
        # populate detection status array
1911
        # 1:detected, 0:missed, -1:below IWA, -2:beyond OWA
1912
        if len(pInds) > 0:
1✔
1913
            detected = (~MD).astype(int)
1✔
1914
            WA = (
1✔
1915
                np.array(
1916
                    [
1917
                        systemParamss[x]["WA"].to_value(u.arcsec)
1918
                        for x in range(len(systemParamss))
1919
                    ]
1920
                )
1921
                << u.arcsec
1922
            )
1923
            detected[np.all(WA < mode["IWA"], 0)] = -1
1✔
1924
            detected[np.all(WA > mode["OWA"], 0)] = -2
1✔
1925

1926
        # if planets are detected, calculate the minimum apparent separation
1927
        smin = None
1✔
1928
        det = detected == 1  # If any of the planets around the star have been detected
1✔
1929
        if np.any(det):
1✔
1930
            smin = np.min(SU.s[pInds[det]])
1✔
1931
            log_det = "   - Detected planet inds %s (%s/%s)" % (
1✔
1932
                pInds[det],
1933
                len(pInds[det]),
1934
                len(pInds),
1935
            )
1936
            self.logger.info(log_det)
1✔
1937
            self.vprint(log_det)
1✔
1938

1939
        # populate the lastDetected array by storing det, JEZ, dMag, and WA
1940
        self.lastDetected[sInd, :] = [
1✔
1941
            det,
1942
            JEZ.flatten(),
1943
            systemParams["dMag"],
1944
            systemParams["WA"].to("arcsec").value,
1945
        ]
1946

1947
        # in case of a FA, generate a random delta mag (between PPro.FAdMag0 and
1948
        # TL.intCutoff_dMag) and working angle (between IWA and min(OWA, a_max))
1949
        if FA:
1✔
1950
            WA = (
×
1951
                np.random.uniform(
1952
                    mode["IWA"].to_value(u.arcsec),
1953
                    np.minimum(
1954
                        mode["OWA"], np.arctan(max(PPop.arange) / TL.dist[sInd])
1955
                    ).to_value(u.arcsec),
1956
                )
1957
                << u.arcsec
1958
            )
1959
            dMag = np.random.uniform(PPro.FAdMag0(WA), TL.intCutoff_dMag)
×
1960
            self.lastDetected[sInd, 0] = np.append(self.lastDetected[sInd, 0], True)
×
1961
            self.lastDetected[sInd, 1] = np.append(
×
1962
                self.lastDetected[sInd, 1], TL.JEZ0[mode["hex"]][sInd].flatten()
1963
            )
1964
            self.lastDetected[sInd, 2] = np.append(self.lastDetected[sInd, 2], dMag)
×
1965
            self.lastDetected[sInd, 3] = np.append(
×
1966
                self.lastDetected[sInd, 3], WA.to_value(u.arcsec)
1967
            )
1968
            sminFA = np.tan(WA) * TL.dist[sInd].to("AU")
×
1969
            smin = np.minimum(smin, sminFA) if smin is not None else sminFA
×
1970
            log_FA = "   - False Alarm (WA=%s, dMag=%s)" % (
×
1971
                np.round(WA, 3),
1972
                np.round(dMag, 1),
1973
            )
1974
            self.logger.info(log_FA)
×
1975
            self.vprint(log_FA)
×
1976

1977
        # Schedule Target Revisit
1978
        self.scheduleRevisit(sInd, smin, det, pInds)
1✔
1979

1980
        if self.make_debug_bird_plots:
1✔
1981
            from tools.obs_plot import obs_plot
×
1982

1983
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, detected)
×
1984

1985
        return detected.astype(int), fZ, JEZ, systemParams, SNR, FA
1✔
1986

1987
    def scheduleRevisit(self, sInd, smin, det, pInds):
1✔
1988
        """A Helper Method for scheduling revisits after observation detection
1989

1990
        Updates self.starRevisit attribute only
1991

1992
        Args:
1993
            sInd (int):
1994
                sInd of the star just detected
1995
            smin (~astropy.units.Quantity):
1996
                minimum separation of the planet to star of planet just detected
1997
            det (~np.ndarray(bool)):
1998
                Detection status of all planets in target system
1999
            pInds (~np.ndarray(int)):
2000
                Indices of planets in the target system
2001

2002
        Returns:
2003
            None
2004

2005
        """
2006
        TK = self.TimeKeeping
1✔
2007
        TL = self.TargetList
1✔
2008
        SU = self.SimulatedUniverse
1✔
2009

2010
        # in both cases (detection or false alarm), schedule a revisit
2011
        # based on minimum separation
2012
        Ms = TL.MsTrue[sInd]
1✔
2013
        if smin is not None:  # smin is None if no planet was detected
1✔
2014
            sp = smin
1✔
2015
            if np.any(det):
1✔
2016
                pInd_smin = pInds[det][np.argmin(SU.s[pInds[det]])]
1✔
2017
                Mp = SU.Mp[pInd_smin]
1✔
2018
            else:
2019
                Mp = SU.Mp.mean()
×
2020
            mu = const.G * (Mp + Ms)
1✔
2021
            T = 2.0 * np.pi * np.sqrt(sp**3.0 / mu)
1✔
2022
            t_rev = TK.currentTimeNorm.copy() + T / 2.0
1✔
2023
        # otherwise, revisit based on average of population semi-major axis and mass
2024
        else:
2025
            sp = SU.s.mean()
1✔
2026
            Mp = SU.Mp.mean()
1✔
2027
            mu = const.G * (Mp + Ms)
1✔
2028
            T = 2.0 * np.pi * np.sqrt(sp**3.0 / mu)
1✔
2029
            t_rev = TK.currentTimeNorm.copy() + 0.75 * T
1✔
2030

2031
        if self.revisit_wait is not None:
1✔
2032
            revisit_wait = self.revisit_wait
1✔
2033
            if np.ndim(revisit_wait) > 0:
1✔
2034
                revisit_wait = revisit_wait[sInd]
1✔
2035
            t_rev = TK.currentTimeNorm.copy() + revisit_wait
1✔
2036
        # finally, populate the revisit list (NOTE: sInd becomes a float)
2037
        revisit = np.array([sInd, float(t_rev.to_value(u.day))])
1✔
2038
        if self.starRevisit.size == 0:  # If starRevisit has nothing in it
1✔
2039
            self.starRevisit = np.array([revisit])  # initialize sterRevisit
1✔
2040
        else:
2041
            revInd = np.where(self.starRevisit[:, 0] == sInd)[
1✔
2042
                0
2043
            ]  # indices of the first column of the starRevisit list containing sInd
2044
            if revInd.size == 0:
1✔
2045
                self.starRevisit = np.vstack((self.starRevisit, revisit))
1✔
2046
            else:
2047
                self.starRevisit[revInd, 1] = revisit[1]
×
2048

2049
    def observation_characterization(self, sInd, mode):
1✔
2050
        """Finds if characterizations are possible and relevant information
2051

2052
        Args:
2053
            sInd (int):
2054
                Integer index of the star of interest
2055
            mode (dict):
2056
                Selected observing mode for characterization
2057

2058
        Returns:
2059
            tuple:
2060
                characterized (list(int)):
2061
                    Characterization status for each planet orbiting the observed
2062
                    target star including False Alarm if any, where 1 is full spectrum,
2063
                    -1 partial spectrum, and 0 not characterized
2064
                fZ (astropy.units.Quantity(numpy.ndarray(float))):
2065
                    Surface brightness of local zodiacal light in units of 1/arcsec2
2066
                JEZ (astropy.units.Quantity(numpy.ndarray(float))):
2067
                    Intensity of exo-zodiacal light in units of photons/s/m2/arcsec
2068
                systemParams (dict):
2069
                    Dictionary of time-dependant planet properties averaged over the
2070
                    duration of the integration
2071
                SNR (numpy.ndarray(float)):
2072
                    Characterization signal-to-noise ratio of the observable planets.
2073
                    Defaults to None.
2074
                intTime (astropy.units.Quantity(numpy.ndarray(float))):
2075
                    Selected star characterization time in units of day.
2076
                    Defaults to None.
2077

2078
        """
2079

2080
        OS = self.OpticalSystem
1✔
2081
        ZL = self.ZodiacalLight
1✔
2082
        TL = self.TargetList
1✔
2083
        SU = self.SimulatedUniverse
1✔
2084
        Obs = self.Observatory
1✔
2085
        TK = self.TimeKeeping
1✔
2086

2087
        # selecting appropriate koMap
2088
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
2089

2090
        # find indices of planets around the target
2091
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2092

2093
        # get the detected status, and check if there was a FA
2094
        det = self.lastDetected[sInd, 0]
1✔
2095
        FA = len(det) == len(pInds) + 1
1✔
2096
        if FA:
1✔
2097
            pIndsDet = np.append(pInds, -1)[det]
×
2098
        else:
2099
            pIndsDet = pInds[det]
1✔
2100

2101
        # initialize outputs, and check if there's anything (planet or FA)
2102
        # to characterize
2103
        characterized = np.zeros(len(det), dtype=int)
1✔
2104
        fZ = 0.0 * self.inv_arcsec2
1✔
2105
        JEZ = 0.0 * self.JEZ_unit
1✔
2106
        # write current system params by default
2107
        systemParams = SU.dump_system_params(sInd)
1✔
2108
        SNR = np.zeros(len(det))
1✔
2109
        intTime = None
1✔
2110
        if len(det) == 0:  # nothing to characterize
1✔
2111
            return characterized, fZ, JEZ, systemParams, SNR, intTime
1✔
2112

2113
        # look for last detected planets that have not been fully characterized
2114
        if not (FA):  # only true planets, no FA
1✔
2115
            tochar = self.fullSpectra[pIndsDet] == 0
1✔
2116
        else:  # mix of planets and a FA
2117
            truePlans = pIndsDet[:-1]
×
2118
            tochar = np.append((self.fullSpectra[truePlans] == 0), True)
×
2119

2120
        # 1/ find spacecraft orbital START position including overhead time,
2121
        # and check keepout angle
2122
        if np.any(tochar):
1✔
2123
            # start times
2124
            startTime = (
1✔
2125
                TK.currentTimeAbs.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2126
            )
2127

2128
            # if we're beyond mission end, bail out
2129
            if startTime >= TK.missionFinishAbs:
1✔
2130
                return characterized, fZ, systemParams, SNR, intTime
×
2131

2132
            startTimeNorm = (
1✔
2133
                TK.currentTimeNorm.copy() + mode["syst"]["ohTime"] + Obs.settlingTime
2134
            )
2135
            # planets to characterize
2136
            koTimeInd = np.where(np.round(startTime.value) - self.koTimes.value == 0)[
1✔
2137
                0
2138
            ][
2139
                0
2140
            ]  # find indice where koTime is startTime[0]
2141
            # wherever koMap is 1, the target is observable
2142
            tochar[tochar] = koMap[sInd][koTimeInd]
1✔
2143

2144
        # 2/ if any planet to characterize, find the characterization times
2145
        # at the detected JEZ, dMag, and WA
2146
        if np.any(tochar):
1✔
2147
            fZ = ZL.fZ(Obs, TL, np.array([sInd], ndmin=1), startTime.reshape(1), mode)[
1✔
2148
                0
2149
            ]
2150
            JEZ = self.lastDetected[sInd, 1][det][tochar]
1✔
2151
            dMag = self.lastDetected[sInd, 2][det][tochar]
1✔
2152
            WA = self.lastDetected[sInd, 3][det][tochar] * u.arcsec
1✔
2153
            intTimes = np.zeros(len(tochar)) << u.day
1✔
2154
            intTimes[tochar] = OS.calc_intTime(TL, sInd, fZ, JEZ, dMag, WA, mode, TK=TK)
1✔
2155
            intTimes[~np.isfinite(intTimes)] = self.zero_d
1✔
2156
            # add a predetermined margin to the integration times
2157
            intTimes = intTimes * (1.0 + self.charMargin)
1✔
2158
            # apply time multiplier
2159
            totTimes = intTimes * (mode["timeMultiplier"])
1✔
2160

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

2165
            # end times
2166
            endTimes = startTime + totTimes
1✔
2167
            endTimesNorm = startTimeNorm + totTimes
1✔
2168
            # planets to characterize
2169
            tochar = (
1✔
2170
                (totTimes > 0)
2171
                & (totTimes <= OS.intCutoff)
2172
                & (endTimesNorm <= TK.OBendTimes[TK.OBnumber])
2173
            )
2174

2175
        # 3/ is target still observable at the end of any char time?
2176
        if np.any(tochar) and Obs.checkKeepoutEnd:
1✔
2177
            koTimeInds = np.zeros(len(endTimes.value[tochar]), dtype=int)
1✔
2178
            # find index in koMap where each endTime is closest to koTimes
2179
            for t, endTime in enumerate(endTimes.value[tochar]):
1✔
2180
                if endTime > self.koTimes.value[-1]:
1✔
2181
                    # case where endTime exceeds largest koTimes element
2182
                    endTimeInBounds = np.where(
×
2183
                        np.floor(endTime) - self.koTimes.value == 0
2184
                    )[0]
2185
                    koTimeInds[t] = (
×
2186
                        endTimeInBounds[0] if endTimeInBounds.size != 0 else -1
2187
                    )
2188
                else:
2189
                    koTimeInds[t] = np.where(
1✔
2190
                        np.round(endTime) - self.koTimes.value == 0
2191
                    )[0][
2192
                        0
2193
                    ]  # find indice where koTime is endTimes[0]
2194
            tochar[tochar] = [koMap[sInd][koT] if koT >= 0 else 0 for koT in koTimeInds]
1✔
2195

2196
        # 4/ if yes, allocate the overhead time, and perform the characterization
2197
        # for the maximum char time
2198
        if np.any(tochar):
1✔
2199
            # Save Current Time before attempting time allocation
2200
            currentTimeNorm = TK.currentTimeNorm.copy()
1✔
2201
            currentTimeAbs = TK.currentTimeAbs.copy()
1✔
2202

2203
            # Allocate Time
2204
            intTime = np.max(intTimes[tochar])
1✔
2205
            extraTime = intTime * (mode["timeMultiplier"] - 1.0)
1✔
2206
            success = TK.allocate_time(
1✔
2207
                intTime + extraTime + mode["syst"]["ohTime"] + Obs.settlingTime, True
2208
            )  # allocates time
2209
            if not (success):  # Time was not successfully allocated
1✔
2210
                char_intTime = None
×
2211
                lenChar = len(pInds) + 1 if FA else len(pInds)
×
2212
                characterized = np.zeros(lenChar, dtype=float)
×
2213
                char_JEZ = np.zeros(lenChar, dtype=float) << self.JEZ_unit
×
2214
                char_SNR = np.zeros(lenChar, dtype=float)
×
2215
                char_fZ = 0.0 * self.inv_arcsec2
×
2216
                char_systemParams = SU.dump_system_params(sInd)
×
2217
                return (
×
2218
                    characterized,
2219
                    char_fZ,
2220
                    char_JEZ,
2221
                    char_systemParams,
2222
                    char_SNR,
2223
                    char_intTime,
2224
                )
2225

2226
            pIndsChar = pIndsDet[tochar]
1✔
2227
            log_char = "   - Charact. planet inds %s (%s/%s detected)" % (
1✔
2228
                pIndsChar,
2229
                len(pIndsChar),
2230
                len(pIndsDet),
2231
            )
2232
            self.logger.info(log_char)
1✔
2233
            self.vprint(log_char)
1✔
2234

2235
            # SNR CALCULATION:
2236
            # first, calculate SNR for observable planets (without false alarm)
2237
            planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
1✔
2238
            SNRplans = np.zeros(len(planinds))
1✔
2239
            if len(planinds) > 0:
1✔
2240
                # initialize arrays for SNR integration
2241
                fZs = np.zeros(self.ntFlux) << self.inv_arcsec2
1✔
2242
                systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
2243
                JEZs = np.zeros((self.ntFlux, len(planinds))) << self.JEZ_unit
1✔
2244
                Ss = np.zeros((self.ntFlux, len(planinds)))
1✔
2245
                Ns = np.zeros((self.ntFlux, len(planinds)))
1✔
2246
                # integrate the signal (planet flux) and noise
2247
                dt = intTime / float(self.ntFlux)
1✔
2248
                timePlus = (
1✔
2249
                    Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2250
                )  # accounts for the time since the current time
2251
                for i in range(self.ntFlux):
1✔
2252
                    # allocate first half of dt
2253
                    timePlus += dt / 2.0
1✔
2254
                    # calculate current zodiacal light brightness
2255
                    fZs[i] = ZL.fZ(
1✔
2256
                        Obs,
2257
                        TL,
2258
                        np.array([sInd], ndmin=1),
2259
                        (currentTimeAbs + timePlus).reshape(1),
2260
                        mode,
2261
                    )[0]
2262
                    # propagate the system to match up with current time
2263
                    SU.propag_system(
1✔
2264
                        sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2265
                    )
2266
                    self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
2267
                    # Calculate the exozodi intensity
2268
                    JEZs[i] = SU.scale_JEZ(sInd, mode, pInds=planinds)
1✔
2269
                    # save planet parameters
2270
                    systemParamss[i] = SU.dump_system_params(sInd)
1✔
2271
                    # calculate signal and noise (electron count rates)
2272
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
2273
                        sInd, planinds, dt, mode, fZ=fZs[i], JEZ=JEZs[i]
2274
                    )
2275
                    # allocate second half of dt
2276
                    timePlus += dt / 2.0
1✔
2277

2278
                # average output parameters
2279
                fZ = np.mean(fZs)
1✔
2280
                JEZ = np.mean(JEZs, axis=0)
1✔
2281
                systemParams = {
1✔
2282
                    key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2283
                    / float(self.ntFlux)
2284
                    for key in sorted(systemParamss[0])
2285
                }
2286
                # calculate planets SNR
2287
                S = Ss.sum(0)
1✔
2288
                N = Ns.sum(0)
1✔
2289
                SNRplans[N > 0] = S[N > 0] / N[N > 0]
1✔
2290
                # allocate extra time for timeMultiplier
2291

2292
            # if only a FA, just save zodiacal brightness in the middle of the
2293
            # integration
2294
            else:
2295
                totTime = intTime * (mode["timeMultiplier"])
×
2296
                fZ = ZL.fZ(
×
2297
                    Obs,
2298
                    TL,
2299
                    np.array([sInd], ndmin=1),
2300
                    (currentTimeAbs + totTime / 2.0).reshape(1),
2301
                    mode,
2302
                )[0]
2303
                # Use the default star value if no planets
2304
                JEZ = TL.JEZ0[mode["hex"]][sInd]
×
2305

2306
            # calculate the false alarm SNR (if any)
2307
            SNRfa = []
1✔
2308
            if pIndsChar[-1] == -1:
1✔
2309
                JEZ = self.lastDetected[sInd, 1][-1]
×
2310
                dMag = self.lastDetected[sInd, 2][-1]
×
2311
                WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
2312
                C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, JEZ, dMag, WA, mode, TK=TK)
×
2313
                S = (C_p * intTime).decompose().value
×
2314
                N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2315
                SNRfa = S / N if N > 0.0 else 0.0
×
2316

2317
            # save all SNRs (planets and FA) to one array
2318
            SNRinds = np.where(det)[0][tochar]
1✔
2319
            SNR[SNRinds] = np.append(SNRplans, SNRfa)
1✔
2320

2321
            # now, store characterization status: 1 for full spectrum,
2322
            # -1 for partial spectrum, 0 for not characterized
2323
            char = SNR >= mode["SNR"]
1✔
2324
            # initialize with full spectra
2325
            characterized = char.astype(int)
1✔
2326
            WAchar = self.lastDetected[sInd, 3][char] * u.arcsec
1✔
2327
            # find the current WAs of characterized planets
2328
            WAs = systemParams["WA"]
1✔
2329
            if FA:
1✔
2330
                WAs = np.append(WAs, self.lastDetected[sInd, 3][-1] * u.arcsec)
×
2331
            # check for partial spectra (for coronagraphs only)
2332
            if not (mode["syst"]["occulter"]):
1✔
2333
                IWA_max = mode["IWA"] * (1.0 + mode["BW"] / 2.0)
1✔
2334
                OWA_min = mode["OWA"] * (1.0 - mode["BW"] / 2.0)
1✔
2335
                char[char] = (WAchar < IWA_max) | (WAchar > OWA_min)
1✔
2336
                characterized[char] = -1
1✔
2337
            # encode results in spectra lists (only for planets, not FA)
2338
            charplans = characterized[:-1] if FA else characterized
1✔
2339
            self.fullSpectra[pInds[charplans == 1]] += 1
1✔
2340
            self.partialSpectra[pInds[charplans == -1]] += 1
1✔
2341

2342
        if self.make_debug_bird_plots:
1✔
2343
            from tools.obs_plot import obs_plot
×
2344

2345
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2346

2347
        return characterized.astype(int), fZ, JEZ, systemParams, SNR, intTime
1✔
2348

2349
    def calc_signal_noise(
1✔
2350
        self, sInd, pInds, t_int, mode, fZ=None, JEZ=None, dMag=None, WA=None
2351
    ):
2352
        """Calculates the signal and noise fluxes for a given time interval. Called
2353
        by observation_detection and observation_characterization methods in the
2354
        SurveySimulation module.
2355

2356
        Args:
2357
            sInd (int):
2358
                Integer index of the star of interest
2359
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2360
                Integration time interval in units of day
2361
            pInds (int):
2362
                Integer indices of the planets of interest
2363
            mode (dict):
2364
                Selected observing mode (from OpticalSystem)
2365
            fZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2366
                Surface brightness of local zodiacal light in units of 1/arcsec2
2367
            JEZ (~astropy.units.Quantity(~numpy.ndarray(float))):
2368
                Intensity of exo-zodiacal light in units of ph/s/m2/arcsec2
2369
            dMag (~numpy.ndarray(float)):
2370
                Differences in magnitude between planets and their host star
2371
            WA (~astropy.units.Quantity(~numpy.ndarray(float))):
2372
                Working angles of the planets of interest in units of arcsec
2373

2374
        Returns:
2375
            tuple:
2376
                Signal (float):
2377
                    Counts of signal
2378
                Noise (float):
2379
                    Counts of background noise variance
2380

2381
        """
2382

2383
        OS = self.OpticalSystem
1✔
2384
        ZL = self.ZodiacalLight
1✔
2385
        TL = self.TargetList
1✔
2386
        SU = self.SimulatedUniverse
1✔
2387
        Obs = self.Observatory
1✔
2388
        TK = self.TimeKeeping
1✔
2389

2390
        # calculate optional parameters if not provided
2391
        fZ = (
1✔
2392
            fZ
2393
            if (fZ is not None)
2394
            else ZL.fZ(
2395
                Obs,
2396
                TL,
2397
                np.array([sInd], ndmin=1),
2398
                TK.currentTimeAbs.copy().reshape(1),
2399
                mode,
2400
            )[0]
2401
        )
2402
        JEZ = (
1✔
2403
            u.Quantity(JEZ, ndmin=1)
2404
            if (JEZ is not None)
2405
            else SU.scale_JEZ(sInd, mode, pInds=pInds)
2406
        )
2407

2408
        # if lucky_planets, use lucky planet params for dMag and WA
2409
        if SU.lucky_planets and mode in list(
1✔
2410
            filter(lambda mode: "spec" in mode["inst"]["name"], OS.observingModes)
2411
        ):
2412
            phi = (1 / np.pi) * np.ones(len(SU.d))
×
2413
            dMag = deltaMag(SU.p, SU.Rp, SU.d, phi)[pInds]  # delta magnitude
×
2414
            # working angle
2415
            WA = (
×
2416
                np.arctan(
2417
                    SU.a.to_value(u.AU)
2418
                    * self.AU2pc
2419
                    / TL.dist[SU.plan2star].to_value(u.pc)
2420
                )
2421
                * self.rad2arcsec
2422
                << u.arcsec
2423
            )[pInds]
2424
        else:
2425
            dMag = dMag if (dMag is not None) else SU.dMag[pInds]
1✔
2426
            WA = WA if (WA is not None) else SU.WA[pInds]
1✔
2427

2428
        # initialize Signal and Noise arrays
2429
        Signal = np.zeros(len(pInds))
1✔
2430
        Noise = np.zeros(len(pInds))
1✔
2431

2432
        # find observable planets wrt IWA-OWA range
2433
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2434

2435
        if np.any(obs):
1✔
2436
            # find electron counts
2437
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(
1✔
2438
                TL, sInd, fZ, JEZ[obs], dMag[obs], WA[obs], mode, TK=TK
2439
            )
2440
            # calculate signal and noise levels (based on Nemati14 formula)
2441
            _t_int_s = t_int.to_value(u.d) * self.day2sec
1✔
2442

2443
            Signal[obs] = C_p.to_value(self.inv_s) * _t_int_s
1✔
2444
            Noise[obs] = np.sqrt(
1✔
2445
                (
2446
                    C_b.to_value(self.inv_s) * _t_int_s
2447
                    + (C_sp.to_value(self.inv_s) * _t_int_s) ** 2
2448
                )
2449
            )
2450

2451
        return Signal, Noise
1✔
2452

2453
    def update_occulter_mass(self, DRM, sInd, t_int, skMode):
1✔
2454
        """Updates the occulter wet mass in the Observatory module, and stores all
2455
        the occulter related values in the DRM array.
2456

2457
        Args:
2458
            DRM (dict):
2459
                Design Reference Mission, contains the results of one complete
2460
                observation (detection and characterization)
2461
            sInd (int):
2462
                Integer index of the star of interest
2463
            t_int (~astropy.units.Quantity(~numpy.ndarray(float))):
2464
                Selected star integration time (for detection or characterization)
2465
                in units of day
2466
            skMode (str):
2467
                Station keeping observing mode type ('det' or 'char')
2468

2469
        Returns:
2470
            dict:
2471
                Design Reference Mission dictionary, contains the results of one
2472
                complete observation
2473

2474
        """
2475

2476
        TL = self.TargetList
×
2477
        Obs = self.Observatory
×
2478
        TK = self.TimeKeeping
×
2479

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

2482
        # decrement mass for station-keeping
2483
        dF_lateral, dF_axial, intMdot, mass_used, deltaV = Obs.mass_dec_sk(
×
2484
            TL, sInd, TK.currentTimeAbs.copy(), t_int
2485
        )
2486

2487
        DRM[skMode + "_dV"] = deltaV.to("m/s")
×
2488
        DRM[skMode + "_mass_used"] = mass_used.to("kg")
×
2489
        DRM[skMode + "_dF_lateral"] = dF_lateral.to("N")
×
2490
        DRM[skMode + "_dF_axial"] = dF_axial.to("N")
×
2491
        # update spacecraft mass
2492
        Obs.scMass = Obs.scMass - mass_used
×
2493
        DRM["scMass"] = Obs.scMass.to("kg")
×
2494
        if Obs.twotanks:
×
2495
            Obs.skMass = Obs.skMass - mass_used
×
2496
            DRM["skMass"] = Obs.skMass.to("kg")
×
2497

2498
        return DRM
×
2499

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

2503
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2504
        objects with their own outspecs.
2505

2506
        Args:
2507
            genNewPlanets (bool):
2508
                Generate all new planets based on the original input specification.
2509
                If False, then the original planets will remain. Setting to True forces
2510
                ``rewindPlanets`` to be True as well. Defaults True.
2511
            rewindPlanets (bool):
2512
                Reset the current set of planets to their original orbital phases.
2513
                If both genNewPlanets and rewindPlanet  are False, the original planets
2514
                will be retained and will not be rewound to their initial starting
2515
                locations (i.e., all systems will remain at the times they were at the
2516
                end of the last run, thereby effectively randomizing planet phases.
2517
                Defaults True.
2518
            seed (int, optional):
2519
                Random seed to use for all random number generation. If None (default)
2520
                a new random seed will be generated when re-initializing the
2521
                SurveySimulation.
2522

2523
        """
2524

2525
        SU = self.SimulatedUniverse
1✔
2526
        TK = self.TimeKeeping
1✔
2527
        TL = self.TargetList
1✔
2528

2529
        # re-initialize SurveySimulation arrays
2530
        specs = self._outspec
1✔
2531
        specs["modules"] = self.modules
1✔
2532

2533
        if seed is None:  # pop the seed so a new one is generated
1✔
2534
            if "seed" in specs:
1✔
2535
                specs.pop("seed")
1✔
2536
        else:  # if seed is provided, replace seed with provided seed
2537
            specs["seed"] = seed
×
2538

2539
        # reset mission time, re-init surveysim and observatory
2540
        TK.__init__(**TK._outspec)
1✔
2541
        self.__init__(**specs)
1✔
2542
        self.Observatory.__init__(**self.Observatory._outspec)
1✔
2543

2544
        # generate new planets if requested (default)
2545
        if genNewPlanets:
1✔
2546
            TL.stellar_mass()
1✔
2547
            TL.I = TL.gen_inclinations(TL.PlanetPopulation.Irange)  # noqa: E741
1✔
2548
            SU.gen_physical_properties(**SU._outspec)
1✔
2549
            rewindPlanets = True
1✔
2550
        # re-initialize systems if requested (default)
2551
        if rewindPlanets:
1✔
2552
            SU.init_systems()
1✔
2553

2554
        # reset helper arrays
2555
        self.initializeStorageArrays()
1✔
2556

2557
        self.vprint("Simulation reset.")
1✔
2558

2559
    def genOutSpec(
1✔
2560
        self,
2561
        tofile: Optional[str] = None,
2562
        starting_outspec: Optional[Dict[str, Any]] = None,
2563
        modnames: bool = False,
2564
    ) -> Dict[str, Any]:
2565
        """Join all _outspec dicts from all modules into one output dict
2566
        and optionally write out to JSON file on disk.
2567

2568
        Args:
2569
            tofile (str, optional):
2570
                Name of the file containing all output specifications (outspecs).
2571
                Defaults to None.
2572
            starting_outspec (dict, optional):
2573
                Initial outspec (from MissionSim). Defaults to None.
2574
            modnames (bool):
2575
                If True, populate outspec dictionary with the module it originated from,
2576
                instead of the actual value of the keyword. Defaults False.
2577

2578
        Returns:
2579
            dict:
2580
                Dictionary containing the full :ref:`sec:inputspec`, including all
2581
                filled-in default values. Combination of all individual module _outspec
2582
                attributes.
2583

2584
        """
2585

2586
        # start with our own outspec
2587
        if modnames:
1✔
2588
            out = copy.copy(self._outspec)
1✔
2589
            for k in out:
1✔
2590
                out[k] = self.__class__.__name__
1✔
2591
        else:
2592
            out = copy.deepcopy(self._outspec)
1✔
2593

2594
        # Add any provided other outspec
2595
        if starting_outspec is not None:
1✔
2596
            out.update(starting_outspec)
1✔
2597

2598
        # add in all modules _outspec's
2599
        for module in self.modules.values():
1✔
2600
            if modnames:
1✔
2601
                tmp = copy.copy(module._outspec)
1✔
2602
                for k in tmp:
1✔
2603
                    tmp[k] = module.__class__.__name__
1✔
2604
            else:
2605
                tmp = module._outspec
1✔
2606
            out.update(tmp)
1✔
2607

2608
        # add in the specific module names used
2609
        out["modules"] = {}
1✔
2610
        for mod_name, module in self.modules.items():
1✔
2611
            # find the module file
2612
            mod_name_full = module.__module__
1✔
2613
            if mod_name_full.startswith("EXOSIMS"):
1✔
2614
                # take just its short name if it is in EXOSIMS
2615
                mod_name_short = mod_name_full.split(".")[-1]
1✔
2616
            else:
2617
                # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2618
                mod_name_short = re.sub(
×
2619
                    r"\.pyc$", ".py", inspect.getfile(module.__class__)
2620
                )
2621
            out["modules"][mod_name] = mod_name_short
1✔
2622
        # add catalog name
2623
        module = self.TargetList.StarCatalog
1✔
2624
        mod_name_full = module.__module__
1✔
2625
        if mod_name_full.startswith("EXOSIMS"):
1✔
2626
            # take just its short name if it is in EXOSIMS
2627
            mod_name_short = mod_name_full.split(".")[-1]
1✔
2628
        else:
2629
            # take its full path if it is not in EXOSIMS - changing .pyc -> .py
2630
            mod_name_short = re.sub(r"\.pyc$", ".py", inspect.getfile(module.__class__))
×
2631
        out["modules"]["StarCatalog"] = mod_name_short
1✔
2632

2633
        # if we don't know about the SurveyEnsemble, just write a blank to the output
2634
        if "SurveyEnsemble" not in out["modules"]:
1✔
2635
            out["modules"]["SurveyEnsemble"] = " "
1✔
2636

2637
        # get version and Git information
2638
        out["Version"] = get_version()
1✔
2639

2640
        # dump to file
2641
        if tofile is not None:
1✔
2642
            with open(tofile, "w") as outfile:
1✔
2643
                json.dump(
1✔
2644
                    out,
2645
                    outfile,
2646
                    sort_keys=True,
2647
                    indent=4,
2648
                    ensure_ascii=False,
2649
                    separators=(",", ": "),
2650
                    default=array_encoder,
2651
                )
2652

2653
        return out
1✔
2654

2655
    def generateHashfName(self, specs):
1✔
2656
        """Generate cached file Hashname
2657

2658
        Requires a .XXX appended to end of hashname for each individual use case
2659

2660
        Args:
2661
            specs (dict):
2662
                :ref:`sec:inputspec`
2663

2664
        Returns:
2665
            str:
2666
                Unique indentifier string for cahce products from this set of modules
2667
                and inputs
2668
        """
2669
        # Allows cachefname to be predefined
2670
        if "cachefname" in specs:
1✔
2671
            return specs["cachefname"]
×
2672

2673
        cachefname = ""  # declares cachefname
1✔
2674
        mods = ["Completeness", "TargetList", "OpticalSystem"]  # modules to look at
1✔
2675
        tmp = (
1✔
2676
            self.Completeness.PlanetPopulation.__class__.__name__
2677
            + self.Completeness.PlanetPhysicalModel.__class__.__name__
2678
            + self.PlanetPopulation.__class__.__name__
2679
            + self.SimulatedUniverse.__class__.__name__
2680
            + self.PlanetPhysicalModel.__class__.__name__
2681
        )
2682

2683
        if "selectionMetric" in specs:
1✔
2684
            tmp += specs["selectionMetric"]
×
2685
        if "Izod" in specs:
1✔
2686
            tmp += specs["Izod"]
×
2687
        if "maxiter" in specs:
1✔
2688
            tmp += str(specs["maxiter"])
×
2689
        if "ftol" in specs:
1✔
2690
            tmp += str(specs["ftol"])
×
2691
        if "missionLife" in specs:
1✔
2692
            tmp += str(specs["missionLife"])
1✔
2693
        if "missionPortion" in specs:
1✔
2694
            tmp += str(specs["missionPortion"])
1✔
2695
        if "smaknee" in specs:
1✔
2696
            tmp += str(specs["smaknee"])
×
2697
        if "koAngleMax" in specs:
1✔
2698
            tmp += str(specs["koAngleMax"])
×
2699
        tmp += str(np.sum(self.Completeness.PlanetPopulation.arange.value))
1✔
2700
        tmp += str(np.sum(self.Completeness.PlanetPopulation.Rprange.value))
1✔
2701
        tmp += str(np.sum(self.Completeness.PlanetPopulation.erange))
1✔
2702
        tmp += str(
1✔
2703
            self.Completeness.PlanetPopulation.PlanetPhysicalModel.whichPlanetPhaseFunction  # noqa: E501
2704
        )
2705
        tmp += str(np.sum(self.PlanetPopulation.arange.value))
1✔
2706
        tmp += str(np.sum(self.PlanetPopulation.Rprange.value))
1✔
2707
        tmp += str(np.sum(self.PlanetPopulation.erange))
1✔
2708
        tmp += str(self.PlanetPopulation.PlanetPhysicalModel.whichPlanetPhaseFunction)
1✔
2709

2710
        for mod in mods:
1✔
2711
            cachefname += self.modules[mod].__module__.split(".")[
1✔
2712
                -1
2713
            ]  # add module name to end of cachefname
2714
        cachefname += hashlib.md5(
1✔
2715
            (str(self.TargetList.Name) + tmp).encode("utf-8")
2716
        ).hexdigest()  # turn cachefname into hashlib
2717
        cachefname = os.path.join(
1✔
2718
            self.cachedir, cachefname + os.extsep
2719
        )  # join into filepath and fname
2720
        # Needs file terminator (.starkt0, .t0, etc) appended done by each individual
2721
        # use case.
2722
        return cachefname
1✔
2723

2724
    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
1✔
2725
        """Helper method for Overloading Revisit Filtering
2726

2727
        Args:
2728
            sInds (~numpy.ndarray(int)):
2729
                Indices of stars still in observation list
2730
            tmpCurrentTimeNorm (astropy.units.Quantity):
2731
                Normalized simulation time
2732

2733
        Returns:
2734
            ~numpy.ndarray(int):
2735
                indices of stars still in observation list
2736
        """
2737

2738
        tovisit = np.zeros(self.TargetList.nStars, dtype=bool)
1✔
2739
        if len(sInds) > 0:
1✔
2740
            # Check that no star has exceeded the number of revisits and the indicies
2741
            # of all considered stars have minimum number of observations
2742
            # This should prevent revisits so long as all stars have not
2743
            # been observed
2744
            tovisit[sInds] = (self.starVisits[sInds] == min(self.starVisits[sInds])) & (
1✔
2745
                self.starVisits[sInds] < self.nVisitsMax
2746
            )
2747
            if self.starRevisit.size != 0:
1✔
2748
                ind_rev = self.revisit_inds(sInds, tmpCurrentTimeNorm)
1✔
2749
                tovisit[ind_rev] = self.starVisits[ind_rev] < self.nVisitsMax
1✔
2750
            sInds = np.where(tovisit)[0]
1✔
2751
        return sInds
1✔
2752

2753
    def is_earthlike(self, plan_inds, sInd):
1✔
2754
        """Is the planet earthlike? Returns boolean array that's True for Earth-like
2755
        planets.
2756

2757

2758
        Args:
2759
            plan_inds(~numpy.ndarray(int)):
2760
                Planet indices
2761
            sInd (int):
2762
                Star index
2763

2764
        Returns:
2765
            ~numpy.ndarray(bool):
2766
                Array of same dimension as plan_inds input that's True for Earthlike
2767
                planets and false otherwise.
2768
        """
2769
        TL = self.TargetList
1✔
2770
        SU = self.SimulatedUniverse
1✔
2771
        PPop = self.PlanetPopulation
1✔
2772

2773
        # extract planet and star properties
2774
        Rp_plan = SU.Rp[plan_inds].value
1✔
2775
        L_star = TL.L[sInd]
1✔
2776
        if PPop.scaleOrbits:
1✔
2777
            a_plan = (SU.a[plan_inds] / np.sqrt(L_star)).value
×
2778
        else:
2779
            a_plan = (SU.a[plan_inds]).value
1✔
2780
        # Definition: planet radius (in earth radii) and solar-equivalent luminosity
2781
        # must be between the given bounds.
2782
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
1✔
2783
        # We use the numpy versions so that plan_ind can be a numpy vector.
2784
        return np.logical_and(
1✔
2785
            np.logical_and(Rp_plan >= Rp_plan_lo, Rp_plan <= 1.4),
2786
            np.logical_and(a_plan >= 0.95, a_plan <= 1.67),
2787
        )
2788

2789
    def find_known_plans(self):
1✔
2790
        """
2791
        Find and return list of known RV stars and list of stars with earthlike planets
2792
        based on info from David Plavchan dated 12/24/2018
2793
        """
2794
        TL = self.TargetList
×
2795
        SU = self.SimulatedUniverse
×
2796
        PPop = self.PlanetPopulation
×
2797
        L_star = TL.L[SU.plan2star]
×
2798

2799
        c = 28.4 * u.m / u.s
×
2800
        Mj = 317.8 * u.earthMass
×
2801
        Mpj = SU.Mp / Mj  # planet masses in jupiter mass units
×
2802
        Ms = TL.MsTrue[SU.plan2star]
×
2803
        Teff = TL.Teff[SU.plan2star]
×
2804
        mu = const.G * (SU.Mp + Ms)
×
2805
        T = (2.0 * np.pi * np.sqrt(SU.a**3 / mu)).to(u.yr)
×
2806
        e = SU.e
×
2807

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

2811
        K = (
×
2812
            (c / np.sqrt(1 - e[t_filt]))
2813
            * Mpj[t_filt]
2814
            * np.sin(SU.I[t_filt])
2815
            * Ms[t_filt] ** (-2 / 3)
2816
            * T[t_filt] ** (-1 / 3)
2817
        )
2818

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

2824
        if PPop.scaleOrbits:
×
2825
            a_plan = (SU.a / np.sqrt(L_star)).value
×
2826
        else:
2827
            a_plan = SU.a.value
×
2828

2829
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
×
2830

2831
        # pinds in habitable zone
2832
        a_filt = k_filt[np.where((a_plan[k_filt] > 0.95) & (a_plan[k_filt] < 1.67))[0]]
×
2833
        # rocky planets
2834
        r_filt = a_filt[
×
2835
            np.where(
2836
                (SU.Rp.value[a_filt] >= Rp_plan_lo[a_filt])
2837
                & (SU.Rp.value[a_filt] < 1.4)
2838
            )[0]
2839
        ]
2840
        self.known_earths = np.union1d(self.known_earths, r_filt).astype(int)
×
2841

2842
        # these are known_rv stars with earths around them
2843
        known_stars = np.unique(SU.plan2star[k_filt])  # these are known_rv stars
×
2844
        known_rocky = np.unique(SU.plan2star[r_filt])
×
2845

2846
        # if include_known_RV, then filter out all other sInds
2847
        if self.include_known_RV is not None:
×
2848
            HIP_sInds = np.where(np.isin(TL.Name, self.include_known_RV))[0]
×
2849
            known_stars = np.intersect1d(HIP_sInds, known_stars)
×
2850
            known_rocky = np.intersect1d(HIP_sInds, known_rocky)
×
2851
        return known_stars.astype(int), known_rocky.astype(int)
×
2852

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

2856
        The observation window (which includes settling and overhead times)
2857
        is a superset of the integration window (in which photons are collected).
2858

2859
        The observation window begins at startTime. The integration window
2860
        begins at startTime + Obs.settlingTime + mode["ohTime"],
2861
        and the integration itself has a duration of intTime.
2862

2863
        Args:
2864
            sInd (int):
2865
                Integer index of the star of interest
2866
            pIndsChar (int numpy.ndarray):
2867
                Observable planets to characterize. Place (-1) at end to put
2868
                False Alarm parameters at end of returned tuples.
2869
            startTime (astropy.units.Quantity):
2870
                Beginning of observation window in units of day.
2871
            intTime (astropy.units.Quantity):
2872
                Selected star characterization integration time in units of day.
2873
            mode (dict):
2874
                Observing mode for the characterization
2875

2876
        Returns:
2877
            tuple:
2878
                SNR (float numpy.ndarray):
2879
                    Characterization signal-to-noise ratio of the observable planets.
2880
                    Defaults to None. [TBD]
2881
                fZ (astropy.units.Quantity):
2882
                    Surface brightness of local zodiacal light in units of 1/arcsec2
2883
                systemParams (dict):
2884
                    Dictionary of time-dependent planet properties averaged over the
2885
                    duration of the integration.
2886

2887
        """
2888

2889
        OS = self.OpticalSystem
1✔
2890
        ZL = self.ZodiacalLight
1✔
2891
        TL = self.TargetList
1✔
2892
        SU = self.SimulatedUniverse
1✔
2893
        Obs = self.Observatory
1✔
2894
        TK = self.TimeKeeping
1✔
2895

2896
        # time at start of integration window
2897
        currentTimeNorm = startTime
1✔
2898
        currentTimeAbs = TK.missionStart + startTime
1✔
2899

2900
        # first, calculate SNR for observable planets (without false alarm)
2901
        planinds = pIndsChar[:-1] if pIndsChar[-1] == -1 else pIndsChar
1✔
2902
        SNRplans = np.zeros(len(planinds))
1✔
2903
        if len(planinds) > 0:
1✔
2904
            # initialize arrays for SNR integration
2905
            fZs = np.zeros(self.ntFlux) << self.fZ_unit
1✔
2906
            systemParamss = np.empty(self.ntFlux, dtype="object")
1✔
2907
            Ss = np.zeros((self.ntFlux, len(planinds)))
1✔
2908
            Ns = np.zeros((self.ntFlux, len(planinds)))
1✔
2909
            # integrate the signal (planet flux) and noise
2910
            dt = intTime / float(self.ntFlux)
1✔
2911
            timePlus = (
1✔
2912
                Obs.settlingTime.copy() + mode["syst"]["ohTime"].copy()
2913
            )  # accounts for the time since the current time
2914
            for i in range(self.ntFlux):
1✔
2915
                # calculate signal and noise (electron count rates)
2916
                if SU.lucky_planets:
1✔
2917
                    fZs[i] = ZL.fZ(
×
2918
                        Obs,
2919
                        TL,
2920
                        np.array([sInd], ndmin=1),
2921
                        currentTimeAbs.reshape(1),
2922
                        mode,
2923
                    )[0]
2924
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
×
2925
                        sInd, planinds, dt, mode, fZ=fZs[i]
2926
                    )
2927
                # allocate first half of dt
2928
                timePlus += dt / 2.0
1✔
2929
                # calculate current zodiacal light brightness
2930
                fZs[i] = ZL.fZ(
1✔
2931
                    Obs,
2932
                    TL,
2933
                    np.array([sInd], ndmin=1),
2934
                    (currentTimeAbs + timePlus).reshape(1),
2935
                    mode,
2936
                )[0]
2937
                # propagate the system to match up with current time
2938
                SU.propag_system(
1✔
2939
                    sInd, currentTimeNorm + timePlus - self.propagTimes[sInd]
2940
                )
2941
                self.propagTimes[sInd] = currentTimeNorm + timePlus
1✔
2942
                # time-varying planet params (WA, dMag, phi, fEZ, d)
2943
                systemParamss[i] = SU.dump_system_params(sInd)
1✔
2944
                # calculate signal and noise (electron count rates)
2945
                if not SU.lucky_planets:
1✔
2946
                    Ss[i, :], Ns[i, :] = self.calc_signal_noise(
1✔
2947
                        sInd, planinds, dt, mode, fZ=fZs[i]
2948
                    )
2949
                # allocate second half of dt
2950
                timePlus += dt / 2.0
1✔
2951

2952
            # average output parameters
2953
            fZ = np.mean(fZs)
1✔
2954
            systemParams = {
1✔
2955
                key: sum([systemParamss[x][key] for x in range(self.ntFlux)])
2956
                / float(self.ntFlux)
2957
                for key in sorted(systemParamss[0])
2958
            }
2959
            # calculate planets SNR
2960
            S = Ss.sum(0)
1✔
2961
            N = Ns.sum(0)
1✔
2962
            SNRplans[N > 0] = S[N > 0] / N[N > 0]
1✔
2963
            # allocate extra time for timeMultiplier
2964

2965
        # if only a FA, just save zodiacal brightness
2966
        # in the middle of the integration
2967
        else:
2968
            totTime = intTime * (mode["timeMultiplier"])
×
2969
            fZ = ZL.fZ(
×
2970
                Obs,
2971
                TL,
2972
                np.array([sInd], ndmin=1),
2973
                (currentTimeAbs.copy() + totTime / 2.0).reshape(1),
2974
                mode,
2975
            )[0]
2976

2977
        # calculate the false alarm SNR (if any)
2978
        SNRfa = []
1✔
2979
        if pIndsChar[-1] == -1:
1✔
2980
            # Note: these attributes may not exist for all schedulers
2981
            fEZ = self.lastDetected[sInd, 1][-1] << self.fZ_unit
×
2982
            dMag = self.lastDetected[sInd, 2][-1]
×
2983
            WA = self.lastDetected[sInd, 3][-1] * u.arcsec
×
2984
            C_p, C_b, C_sp = OS.Cp_Cb_Csp(TL, sInd, fZ, fEZ, dMag, WA, mode)
×
2985
            S = (C_p * intTime).decompose().value
×
2986
            N = np.sqrt((C_b * intTime + (C_sp * intTime) ** 2.0).decompose().value)
×
2987
            SNRfa = S / N if N > 0.0 else 0.0
×
2988

2989
        # SNR vector is of length (#planets), plus 1 if FA
2990
        # This routine leaves SNR = 0 if unknown or not found
2991
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2992
        FA_present = pIndsChar[-1] == -1
1✔
2993
        # there will be one SNR for each entry of pInds_FA
2994
        pInds_FA = np.append(pInds, [-1] if FA_present else np.zeros(0, dtype=int))
1✔
2995

2996
        # boolean index vector into SNR
2997
        #   True iff we have computed a SNR for that planet
2998
        #   False iff we didn't find an SNR (and will leave 0 there)
2999
        #   if FA_present, SNR_plug_in[-1] = True
3000
        SNR_plug_in = np.isin(pInds_FA, pIndsChar)
1✔
3001

3002
        # save all SNRs (planets and FA) to one array
3003
        SNR = np.zeros(len(pInds_FA))
1✔
3004
        # plug in the SNR's we computed (pIndsChar and optional FA)
3005
        SNR[SNR_plug_in] = np.append(SNRplans, SNRfa)
1✔
3006

3007
        return SNR, fZ, systemParams
1✔
3008

3009
    def revisit_inds(self, sInds, tmpCurrentTimeNorm):
1✔
3010
        """Return subset of star indices that are scheduled for revisit.
3011

3012
        Args:
3013
            sInds (~numpy.ndarray(int)):
3014
                Indices of stars to consider
3015
            tmpCurrentTimeNorm (astropy.units.Quantity):
3016
                Normalized simulation time
3017

3018
        Returns:
3019
            ~numpy.ndarray(int):
3020
                Indices of stars whose revisit is scheduled for within `self.dt_max` of
3021
                the current time
3022

3023
        """
3024
        dt_rev = np.abs(self.starRevisit[:, 1] * u.day - tmpCurrentTimeNorm)
1✔
3025
        ind_rev = [
1✔
3026
            int(x) for x in self.starRevisit[dt_rev < self.dt_max, 0] if x in sInds
3027
        ]
3028

3029
        return ind_rev
1✔
3030

3031
    def keepout_filter(self, sInds, startTimes, endTimes, koMap):
1✔
3032
        """Filters stars not observable due to keepout
3033

3034
        Args:
3035
            sInds (~numpy.ndarray(int)):
3036
                indices of stars still in observation list
3037
            startTimes (~numpy.ndarray(float)):
3038
                start times of observations
3039
            endTimes (~numpy.ndarray(float)):
3040
                end times of observations
3041
            koMap (~numpy.ndarray(bool)):
3042
                map where True is for a target unobstructed and observable,
3043
                False is for a target obstructed and unobservable due to keepout zone
3044

3045
        Returns:
3046
            ~numpy.ndarray(int):
3047
                sInds - filtered indices of stars still in observation list
3048

3049
        """
3050
        # find the indices of keepout times that pertain to the start and end times of
3051
        # observation
3052
        startInds = np.searchsorted(self.koTimes.value, startTimes)
×
3053
        endInds = np.searchsorted(self.koTimes.value, endTimes)
×
3054

3055
        # boolean array of available targets (unobstructed between observation start and
3056
        # end time) we include a -1 in the start and +1 in the end to include keepout
3057
        # days enclosing start and end times
3058
        availableTargets = np.array(
×
3059
            [
3060
                np.all(
3061
                    koMap[
3062
                        sInd,
3063
                        max(startInds[sInd] - 1, 0) : min(
3064
                            endInds[sInd] + 1, len(self.koTimes.value) + 1
3065
                        ),
3066
                    ]
3067
                )
3068
                for sInd in sInds
3069
            ],
3070
            dtype="bool",
3071
        )
3072

3073
        return sInds[availableTargets]
×
3074

3075

3076
def array_encoder(obj):
1✔
3077
    r"""Encodes numpy arrays, astropy Times, and astropy Quantities, into JSON.
3078

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

3085
    Args:
3086
        obj (Any):
3087
            Object to encode.
3088

3089
        Returns:
3090
            Any:
3091
                Encoded object
3092

3093
    """
3094

3095
    from astropy.coordinates import SkyCoord
1✔
3096
    from astropy.time import Time
1✔
3097

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

© 2026 Coveralls, Inc