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

dsavransky / EXOSIMS / 21831994050

09 Feb 2026 03:51PM UTC coverage: 65.64% (-0.02%) from 65.656%
21831994050

Pull #457

github

web-flow
Merge 20c9d1a4e into 663fe0e23
Pull Request #457: YAML Input Script Capability

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

9 existing lines in 2 files now uncovered.

9846 of 15000 relevant lines covered (65.64%)

0.66 hits per line

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

62.41
/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
            t_rev = TK.currentTimeNorm.copy() + self.revisit_wait
×
2033
        # finally, populate the revisit list (NOTE: sInd becomes a float)
2034
        revisit = np.array([sInd, t_rev.to_value(u.day)])
1✔
2035
        if self.starRevisit.size == 0:  # If starRevisit has nothing in it
1✔
2036
            self.starRevisit = np.array([revisit])  # initialize sterRevisit
1✔
2037
        else:
2038
            revInd = np.where(self.starRevisit[:, 0] == sInd)[
1✔
2039
                0
2040
            ]  # indices of the first column of the starRevisit list containing sInd
2041
            if revInd.size == 0:
1✔
2042
                self.starRevisit = np.vstack((self.starRevisit, revisit))
1✔
2043
            else:
2044
                self.starRevisit[revInd, 1] = revisit[1]
1✔
2045

2046
    def observation_characterization(self, sInd, mode):
1✔
2047
        """Finds if characterizations are possible and relevant information
2048

2049
        Args:
2050
            sInd (int):
2051
                Integer index of the star of interest
2052
            mode (dict):
2053
                Selected observing mode for characterization
2054

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

2075
        """
2076

2077
        OS = self.OpticalSystem
1✔
2078
        ZL = self.ZodiacalLight
1✔
2079
        TL = self.TargetList
1✔
2080
        SU = self.SimulatedUniverse
1✔
2081
        Obs = self.Observatory
1✔
2082
        TK = self.TimeKeeping
1✔
2083

2084
        # selecting appropriate koMap
2085
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
2086

2087
        # find indices of planets around the target
2088
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2089

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2339
        if self.make_debug_bird_plots:
1✔
2340
            from tools.obs_plot import obs_plot
×
2341

2342
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2343

2344
        return characterized.astype(int), fZ, JEZ, systemParams, SNR, intTime
1✔
2345

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

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

2371
        Returns:
2372
            tuple:
2373
                Signal (float):
2374
                    Counts of signal
2375
                Noise (float):
2376
                    Counts of background noise variance
2377

2378
        """
2379

2380
        OS = self.OpticalSystem
1✔
2381
        ZL = self.ZodiacalLight
1✔
2382
        TL = self.TargetList
1✔
2383
        SU = self.SimulatedUniverse
1✔
2384
        Obs = self.Observatory
1✔
2385
        TK = self.TimeKeeping
1✔
2386

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

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

2425
        # initialize Signal and Noise arrays
2426
        Signal = np.zeros(len(pInds))
1✔
2427
        Noise = np.zeros(len(pInds))
1✔
2428

2429
        # find observable planets wrt IWA-OWA range
2430
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2431

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

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

2448
        return Signal, Noise
1✔
2449

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

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

2466
        Returns:
2467
            dict:
2468
                Design Reference Mission dictionary, contains the results of one
2469
                complete observation
2470

2471
        """
2472

2473
        TL = self.TargetList
×
2474
        Obs = self.Observatory
×
2475
        TK = self.TimeKeeping
×
2476

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

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

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

2495
        return DRM
×
2496

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

2500
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2501
        objects with their own outspecs.
2502

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

2520
        """
2521

2522
        SU = self.SimulatedUniverse
1✔
2523
        TK = self.TimeKeeping
1✔
2524
        TL = self.TargetList
1✔
2525

2526
        # re-initialize SurveySimulation arrays
2527
        specs = self._outspec
1✔
2528
        specs["modules"] = self.modules
1✔
2529

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

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

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

2551
        # reset helper arrays
2552
        self.initializeStorageArrays()
1✔
2553

2554
        self.vprint("Simulation reset.")
1✔
2555

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

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

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

2581
        """
2582

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

2591
        # Add any provided other outspec
2592
        if starting_outspec is not None:
1✔
2593
            out.update(starting_outspec)
1✔
2594

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

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

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

2634
        # get version and Git information
2635
        out["Version"] = get_version()
1✔
2636

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

2650
        return out
1✔
2651

2652
    def generateHashfName(self, specs):
1✔
2653
        """Generate cached file Hashname
2654

2655
        Requires a .XXX appended to end of hashname for each individual use case
2656

2657
        Args:
2658
            specs (dict):
2659
                :ref:`sec:inputspec`
2660

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

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

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

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

2721
    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
1✔
2722
        """Helper method for Overloading Revisit Filtering
2723

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

2730
        Returns:
2731
            ~numpy.ndarray(int):
2732
                indices of stars still in observation list
2733
        """
2734

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

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

2754

2755
        Args:
2756
            plan_inds(~numpy.ndarray(int)):
2757
                Planet indices
2758
            sInd (int):
2759
                Star index
2760

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

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

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

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

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

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

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

2821
        if PPop.scaleOrbits:
×
2822
            a_plan = (SU.a / np.sqrt(L_star)).value
×
2823
        else:
2824
            a_plan = SU.a.value
×
2825

2826
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
×
2827

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

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

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

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

2853
        The observation window (which includes settling and overhead times)
2854
        is a superset of the integration window (in which photons are collected).
2855

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

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

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

2884
        """
2885

2886
        OS = self.OpticalSystem
×
2887
        ZL = self.ZodiacalLight
×
2888
        TL = self.TargetList
×
2889
        SU = self.SimulatedUniverse
×
2890
        Obs = self.Observatory
×
2891
        TK = self.TimeKeeping
×
2892

2893
        # time at start of integration window
2894
        currentTimeNorm = startTime
×
2895
        currentTimeAbs = TK.missionStart + startTime
×
2896

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

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

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

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

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

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

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

3004
        return SNR, fZ, systemParams
×
3005

3006
    def revisit_inds(self, sInds, tmpCurrentTimeNorm):
1✔
3007
        """Return subset of star indices that are scheduled for revisit.
3008

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

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

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

3026
        return ind_rev
1✔
3027

3028
    def keepout_filter(self, sInds, startTimes, endTimes, koMap):
1✔
3029
        """Filters stars not observable due to keepout
3030

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

3042
        Returns:
3043
            ~numpy.ndarray(int):
3044
                sInds - filtered indices of stars still in observation list
3045

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

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

3070
        return sInds[availableTargets]
×
3071

3072

3073
def array_encoder(obj):
1✔
3074
    r"""Encodes numpy arrays, astropy Times, and astropy Quantities, into JSON.
3075

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

3082
    Args:
3083
        obj (Any):
3084
            Object to encode.
3085

3086
        Returns:
3087
            Any:
3088
                Encoded object
3089

3090
    """
3091

3092
    from astropy.coordinates import SkyCoord
1✔
3093
    from astropy.time import Time
1✔
3094

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

© 2026 Coveralls, Inc