Creates a grid of orbits and times pyoorb propagations.

pyoorbTimer.py

#!/usr/bin/python
"""
[email protected]
2019/07/15

Creates a grid of orbits and propagates them 10 days from the defined epoch. 

Installation (using a dedicated conda environment): 
```
conda create -n pyoorb_py36 -c defaults -conda-forge python=3.6 numpy
conda activate pyoorb_py36
conda install -c defaults -c conda-forge openorb
```

USAGE
-----
```python pyoorbTimer.py```
Time to setup pyoorb: 0.22381305694580078 seconds
Time to propagate 1000000 orbits: 295.9671907424927 seconds

```python pyoorbTimer.py --bins 3```
Time to setup pyoorb: 0.18677186965942383 seconds
Time to propagate 729 orbits: 0.18398332595825195 seconds

```python pyoorbTimer.py --bins 4```
Time to setup pyoorb: 0.13042092323303223 seconds
Time to propagate 4096 orbits: 1.001237154006958 seconds

TODO
----
- Check pyoorb outputs for errors [partially done]
- Add memory consumption tracking
- Clean up kwargs
- This script should probably be versioned and stored on GitHub.

"""
import os
import time
import warnings
import argparse
import numpy as np
import pyoorb as oo

ELEMENT_RANGES = {
    # name : [low, high]
    "a_au" : [0.1, 25.0],
    "e" : [0, 1],
    "i_rad" : [0, np.pi],
    "ascNode_rad" : [0, 2*np.pi],
    "argPeri_rad" : [0, 2*np.pi],
    "meanAnom_rad" : [0, 2*np.pi],
    }
BINS = 10


def setupPyoorb():
    time_start = time.time()
    if os.environ.get("OORB_DATA") == None:
        os.environ["OORB_DATA"] = os.path.join(os.environ["CONDA_PREFIX"], "share/openorb")

    # Prepare pyoorb
    ephfile = os.path.join(os.getenv('OORB_DATA'), 'de430.dat')
    oo.pyoorb.oorb_init(ephfile)
    time_end = time.time()

    print("Time to setup pyoorb: {} seconds".format(time_end - time_start))

def testPropagation(bins=BINS, elementRanges=ELEMENT_RANGES):
    # Should consider making numOrbits the kwarg. Although this setup avoids taking a 6-th root.
    numOrbits = bins**6
    
    # Make uniform bins between low and high values for the six orbital elements
    a_bin = np.random.uniform(*elementRanges["a_au"], size=bins)
    e_bin = np.random.uniform(*elementRanges["e"], size=bins)
    i_bin = np.random.uniform(*elementRanges["i_rad"], size=bins)
    ascNode_bin = np.random.uniform(*elementRanges["ascNode_rad"], size=bins)
    argPeri_bin = np.random.uniform(*elementRanges["argPeri_rad"], size=bins)
    meanAnom_bin = np.random.uniform(*elementRanges["meanAnom_rad"], size=bins)

    # Make grid of orbits
    a, e, i, ascNode_rad, argPeri_rad, meanAnom_rad = np.meshgrid(a_bin, e_bin, i_bin, ascNode_bin, argPeri_bin, meanAnom_bin)
    a = a.flatten()
    e = e.flatten()
    i = i.flatten()
    ascNode_rad = ascNode_rad.flatten()
    argPeri_rad = argPeri_rad.flatten()
    meanAnom_rad = meanAnom_rad.flatten()
    orbit_ids = np.arange(0, numOrbits)
    orbitType = [3 for i in range(numOrbits)]
    epoch = [59580 for i in range(numOrbits)]
    timeScale = [1 for i in range(numOrbits)]
    absoluteMagnitude = [18 for i in range(numOrbits)]
    slopeParameter = [0.15 for i in range(numOrbits)]

    # Format orbits in pyoorb "style"
    orbits = np.array(
        np.array([
            orbit_ids,
            a, 
            e, 
            i, 
            ascNode_rad, 
            argPeri_rad,
            meanAnom_rad,
            orbitType,
            epoch,
            timeScale,
            absoluteMagnitude,
            slopeParameter]).T,
         dtype=np.double, 
         order='F')

    # Propagate and wait?
    time_start = time.time()
    epoch_end = np.array([epoch[0] + 10, 1], 
        dtype=np.double, 
        order='F')
    orb, err = oo.pyoorb.oorb_propagation(in_orbits=orbits,
                                          in_epoch=epoch_end,
                                          in_dynmodel='N')
    if np.any(err) == 1:
        # this could be swapped to a failure case if we wanted to make this a unit test
        warnings.warn("""Pyoorb propagation has raised an error.""")

    time_end = time.time()
    print("Time to propagate {} orbits: {} seconds".format(numOrbits, time_end - time_start))
    return


if __name__ == "__main__":

    # Parse arguments... currently only number of bins per element is implemented.
    parser = argparse.ArgumentParser(description="""Creates a grid of orbits and propagates them 10 days from the defined epoch.""")
    parser.add_argument('--bins',
                        default=BINS, 
                        type=int, 
                        help='Number of bins per Keplerian orbital element. Total number of orbits will be bins**6. [Default: {} bins]'.format(BINS))

    args = parser.parse_args()

    setupPyoorb()

    testPropagation(bins=args.bins)