Merge pull request #36 from CalebLammers/caleb

Giant impact emulator example

README.md

@@ -18,8 +18,7 @@ All estimators use a common API to facilitate comparisons between them and with
 
 # Quickstart
 
-Let's predict the probability that a given 3-planet system is stable
-past 1 billion orbits with the XGBoost-based classifier of [Tamayo et al., 2020](http://arxiv.org/abs/2007.06521).
+Let's predict the probability that a given 3-planet system is stable past 1 billion orbits with the XGBoost-based classifier of [Tamayo et al., 2020](http://arxiv.org/abs/2007.06521).
 
 ```python
 import rebound
@@ -37,7 +36,7 @@ print(feature_model.predict_stable(sim))
 # >>> 0.06591137
 ```
 
-That model provides a simple scalar probability of stability over a billion orbits.
+This model provides a simple scalar probability of stability over a billion orbits.
 We can instead estimate its median expected instability time using the deep regressor from [Cranmer et al., 2021](https://arxiv.org/abs/2101.04117).
 
 ```python
@@ -58,7 +57,7 @@ print(median)
 The returned time is expressed in the time units used in setting up the REBOUND Simulation above.
 Since we set the innermost planet orbit to unity, this corresponds to 242570 innermost planet orbits.
 
-Finally, we can compare these results to the semi-analytic criterion of [Tamayo et al., 2021](https://arxiv.org/abs/2106.14863) for how likely the configuration is to be dynamically chaotic.
+We can compare these results to the semi-analytic criterion of [Tamayo et al., 2021](https://arxiv.org/abs/2106.14863) for how likely the configuration is to be dynamically chaotic.
 This is not a one-to-one comparison, but configurations that are chaotic through two-body MMR overlap are generally unstable on long timescales (see paper and examples).
 
 ```python
@@ -71,7 +70,49 @@ print(analytical_model.predict_stable(sim))
 
 To match up with the above classifiers, the analytical classifier returns the probability the configuration is *regular*, i.e., not chaotic.
 A probability of zero therefore corresponds to confidently chaotic.
-See [this example](https://github.com/dtamayo/spock/blob/master/jupyter_examples/QuickStart.ipynb) for more information about the analytical model.
+
+We can also predict the collisional outcome of this system using the MLP model from [Lammers et al., 2024](https://arxiv.org/abs/???).
+
+```python
+from spock import CollisionMergerClassifier
+class_model = CollisionMergerClassifier()
+
+prob_12, prob_23, prob_13 = class_model.predict_collision_probs(sim)
+
+print(prob_12, prob_23, prob_13)
+# >>> 0.2738345 0.49277353 0.23339202
+```
+
+This model returns the probability of a physical collision occurring between planets 1 & 2, 2 & 3, and 1 & 3 when provided a three-planet system. In this case, the instability will most likely result in a collision between planets 2 & 3, but all three outcomes are possible (ejections are exceedingly rare in compact multiplanet systems).
+
+Additionally, we can predict the states of the post-collision planets using the model from [Lammers et al., 2024](https://arxiv.org/abs/???).
+
+```python
+from spock import CollisionOrbitalOutcomeRegressor
+reg_model = CollisionOrbitalOutcomeRegressor()
+
+new_sim = reg_model.predict_collision_outcome(sim, collision_inds=[2, 3])
+
+print(new_sim)
+# >>> <rebound.simulation.Simulation object at 0x303ed70d0, N=3, t=0.0>
+```
+
+Note that this model makes the usual assumption that mergers are perfectly inelastic.
+
+These models are conveniently combined into a giant impact emulator, which predicts instability times, predicts and samples collision pair probabilities, and then handles the collisions. This is repeated until SPOCK predicts the system to be stable on a user-specified timescale (the default is a billion orbits).
+
+```python
+from spock import GiantImpactPhaseEmulator
+emulator = GiantImpactPhaseEmulator()
+
+new_sim = emulator.predict(sim)
+
+print(new_sim)
+# >>> <rebound.simulation.Simulation object at 0x303f05c50, N=3, t=999999999.9999993>
+```
+
+Only one collision takes place in this example system, but the typical use case involves starting with more bodies (~10).
+See [this example](https://github.com/dtamayo/spock/blob/master/jupyter_examples/QuickStart.ipynb) for additional information about the models included in SPOCK, and see [jupyter\_examples/](https://github.com/dtamayo/spock/tree/master/jupyter_examples) for more thorough example applications.
 
 # Examples
 

spock/giant_impact_phase_emulator.py

@@ -40,18 +40,23 @@ def predict(self, sims, tmaxs=None, verbose=False, deepregressor_kwargs={'sample
 
         """
         if isinstance(sims, rb.Simulation): sims = [sims] # passed a single sim
+
+        if verbose: tot_start = time.time() # record start time
 
         # main loop
         sims, tmaxs = self._make_lists(sims, tmaxs)
         while np.any([sim.t < tmaxs[i] for i, sim in enumerate(sims)]): # take another step if any sims are still at t < tmax
             sims = self.step(sims, tmaxs, verbose=verbose, deepregressor_kwargs=deepregressor_kwargs)
             if isinstance(sims, rb.Simulation): sims = [sims] # passed a single sim
 
+        # print total time
+        if verbose:
+            tot_end = time.time()
+            print('Total time:', tot_end - tot_start, 's')
+
         if len(sims) == 1:
             return sims[0] # return single sim
-        else: 
-            return sims
-
+
         return sims
 
     def step(self, sims, tmaxs, verbose=False, deepregressor_kwargs={'samples':100, 'max_model_samples':10}):
@@ -84,6 +89,7 @@ def step(self, sims, tmaxs, verbose=False, deepregressor_kwargs={'samples':100,
             return sims
 
         if verbose:
+            print('Number of sims to update:', len(sims_to_update), '\n')
             print('Predicting trio instability times')
             start = time.time()
 
@@ -194,7 +200,7 @@ def _make_lists(self, sims, tmaxs):
         sims = remove_ejected_ps(sims) # remove ejected/hyperbolic particles (do here so we don't use a negative period for tmaxs)
 
         # use passed value
-        if tmaxs:
+        if not tmaxs is None:
             try:
                 len(tmaxs) == len(sims)
             except: