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

dsavransky / EXOSIMS / 18139527206

30 Sep 2025 06:19PM UTC coverage: 65.5% (-0.1%) from 65.647%
18139527206

push

github

web-flow
Merge pull request #439 from dsavransky/pre-commit-ci-update-config

[pre-commit.ci] pre-commit autoupdate

9768 of 14913 relevant lines covered (65.5%)

0.65 hits per line

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

62.59
/EXOSIMS/Prototypes/SurveySimulation.py
1
# -*- coding: utf-8 -*-
2
import copy
1✔
3
import hashlib
1✔
4
import inspect
1✔
5
import json
1✔
6
import logging
1✔
7
import os
1✔
8
import random as py_random
1✔
9
import re
1✔
10
import sys
1✔
11
import time
1✔
12
from pathlib import Path
1✔
13
from typing import Any, Dict, Optional
1✔
14

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

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

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

29

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

33
    Args:
34

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

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

188
    """
189

190
    _modtype = "SurveySimulation"
1✔
191

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

555
        """
556

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

824
        Returns:
825
            None
826
        """
827

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1055
        .. note::
1056

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

1059
        """
1060

1061
        TL = self.TargetList
1✔
1062

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

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

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

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

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

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

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

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

1130
        return intTimes
1✔
1131

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

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

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

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

1157
        """
1158

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1284
        return sInds, slewTimes, intTimes, dV
×
1285

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1707
            return sInds, intTime, slewTime, dV
×
1708

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

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

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

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

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

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

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

1756
        return sInds, intTime, slewTime, dV
×
1757

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

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

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

1788
        """
1789

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1969
        Updates self.starRevisit attribute only
1970

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

1981
        Returns:
1982
            None
1983

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

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

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

2025
    def observation_characterization(self, sInd, mode):
1✔
2026
        """Finds if characterizations are possible and relevant information
2027

2028
        Args:
2029
            sInd (int):
2030
                Integer index of the star of interest
2031
            mode (dict):
2032
                Selected observing mode for characterization
2033

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

2054
        """
2055

2056
        OS = self.OpticalSystem
1✔
2057
        ZL = self.ZodiacalLight
1✔
2058
        TL = self.TargetList
1✔
2059
        SU = self.SimulatedUniverse
1✔
2060
        Obs = self.Observatory
1✔
2061
        TK = self.TimeKeeping
1✔
2062

2063
        # selecting appropriate koMap
2064
        koMap = self.koMaps[mode["syst"]["name"]]
1✔
2065

2066
        # find indices of planets around the target
2067
        pInds = np.where(SU.plan2star == sInd)[0]
1✔
2068

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2318
        if self.make_debug_bird_plots:
1✔
2319
            from tools.obs_plot import obs_plot
×
2320

2321
            obs_plot(self, systemParams, mode, sInd, pInds, SNR, characterized)
×
2322

2323
        return characterized.astype(int), fZ, JEZ, systemParams, SNR, intTime
1✔
2324

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

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

2350
        Returns:
2351
            tuple:
2352
                Signal (float):
2353
                    Counts of signal
2354
                Noise (float):
2355
                    Counts of background noise variance
2356

2357
        """
2358

2359
        OS = self.OpticalSystem
1✔
2360
        ZL = self.ZodiacalLight
1✔
2361
        TL = self.TargetList
1✔
2362
        SU = self.SimulatedUniverse
1✔
2363
        Obs = self.Observatory
1✔
2364
        TK = self.TimeKeeping
1✔
2365

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

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

2404
        # initialize Signal and Noise arrays
2405
        Signal = np.zeros(len(pInds))
1✔
2406
        Noise = np.zeros(len(pInds))
1✔
2407

2408
        # find observable planets wrt IWA-OWA range
2409
        obs = (WA > mode["IWA"]) & (WA < mode["OWA"])
1✔
2410

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

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

2427
        return Signal, Noise
1✔
2428

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

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

2445
        Returns:
2446
            dict:
2447
                Design Reference Mission dictionary, contains the results of one
2448
                complete observation
2449

2450
        """
2451

2452
        TL = self.TargetList
×
2453
        Obs = self.Observatory
×
2454
        TK = self.TimeKeeping
×
2455

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

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

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

2474
        return DRM
×
2475

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

2479
        This will reinitialize the TimeKeeping, Observatory, and SurveySimulation
2480
        objects with their own outspecs.
2481

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

2499
        """
2500

2501
        SU = self.SimulatedUniverse
1✔
2502
        TK = self.TimeKeeping
1✔
2503
        TL = self.TargetList
1✔
2504

2505
        # re-initialize SurveySimulation arrays
2506
        specs = self._outspec
1✔
2507
        specs["modules"] = self.modules
1✔
2508

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

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

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

2530
        # reset helper arrays
2531
        self.initializeStorageArrays()
1✔
2532

2533
        self.vprint("Simulation reset.")
1✔
2534

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

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

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

2560
        """
2561

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

2570
        # Add any provided other outspec
2571
        if starting_outspec is not None:
1✔
2572
            out.update(starting_outspec)
1✔
2573

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

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

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

2613
        # get version and Git information
2614
        out["Version"] = get_version()
1✔
2615

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

2629
        return out
1✔
2630

2631
    def generateHashfName(self, specs):
1✔
2632
        """Generate cached file Hashname
2633

2634
        Requires a .XXX appended to end of hashname for each individual use case
2635

2636
        Args:
2637
            specs (dict):
2638
                :ref:`sec:inputspec`
2639

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

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

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

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

2700
    def revisitFilter(self, sInds, tmpCurrentTimeNorm):
1✔
2701
        """Helper method for Overloading Revisit Filtering
2702

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

2709
        Returns:
2710
            ~numpy.ndarray(int):
2711
                indices of stars still in observation list
2712
        """
2713

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

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

2733

2734
        Args:
2735
            plan_inds(~numpy.ndarray(int)):
2736
                Planet indices
2737
            sInd (int):
2738
                Star index
2739

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

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

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

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

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

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

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

2800
        if PPop.scaleOrbits:
×
2801
            a_plan = (SU.a / np.sqrt(L_star)).value
×
2802
        else:
2803
            a_plan = SU.a.value
×
2804

2805
        Rp_plan_lo = 0.80 / np.sqrt(a_plan)
×
2806

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

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

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

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

2832
        The observation window (which includes settling and overhead times)
2833
        is a superset of the integration window (in which photons are collected).
2834

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

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

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

2863
        """
2864

2865
        OS = self.OpticalSystem
×
2866
        ZL = self.ZodiacalLight
×
2867
        TL = self.TargetList
×
2868
        SU = self.SimulatedUniverse
×
2869
        Obs = self.Observatory
×
2870
        TK = self.TimeKeeping
×
2871

2872
        # time at start of integration window
2873
        currentTimeNorm = startTime
×
2874
        currentTimeAbs = TK.missionStart + startTime
×
2875

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

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

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

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

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

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

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

2983
        return SNR, fZ, systemParams
×
2984

2985
    def revisit_inds(self, sInds, tmpCurrentTimeNorm):
1✔
2986
        """Return subset of star indices that are scheduled for revisit.
2987

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

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

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

3005
        return ind_rev
1✔
3006

3007
    def keepout_filter(self, sInds, startTimes, endTimes, koMap):
1✔
3008
        """Filters stars not observable due to keepout
3009

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

3021
        Returns:
3022
            ~numpy.ndarray(int):
3023
                sInds - filtered indices of stars still in observation list
3024

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

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

3049
        return sInds[availableTargets]
×
3050

3051

3052
def array_encoder(obj):
1✔
3053
    r"""Encodes numpy arrays, astropy Times, and astropy Quantities, into JSON.
3054

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

3061
    Args:
3062
        obj (Any):
3063
            Object to encode.
3064

3065
        Returns:
3066
            Any:
3067
                Encoded object
3068

3069
    """
3070

3071
    from astropy.coordinates import SkyCoord
1✔
3072
    from astropy.time import Time
1✔
3073

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