Merge pull request #27 from dtamayo/mc/safetensors

Remove PyTorch Lightning dependency

.gitignore

@@ -4,6 +4,7 @@
 __pycache__/
 */.ipynb_checkpoints/
 spock.egg-info/
+build
 training_data/
 data*
 paper_plots/*.csv

README.md

@@ -41,18 +41,24 @@ That model provides a simple scalar probability of stability over a billion orbi
 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
+import numpy as np
 from spock import DeepRegressor
 deep_model = DeepRegressor()
 
-median, lower, upper = deep_model.predict_instability_time(sim, samples=10000)
-print(int(median))
-# >>> 242570.1378387966
+median, lower, upper, samples = deep_model.predict_instability_time(
+    sim, samples=10000, return_samples=True, seed=0
+)
+print(10**np.average(np.log10(samples)))  # Expectation of log-normal
+# >>> 414208.4307974086
+
+print(median)
+# >>> 223792.38826507595
 ```
 
 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. .
+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.
 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

pyproject.toml

@@ -29,27 +29,36 @@ dependencies = [
     "rebound>=3.14.0",           # classifier regressor analytical
     "scikit-learn",              # classifier regressor
     "xgboost>=1.1.0",            # classifier 
-    "matplotlib",                #            regressor
-    "pytorch_lightning>=1.0.0",  #            regressor
     "torch>=1.5.1",              #            regressor
+    "safetensors>=0.4.0,<0.5.0", #            regressor
     "scipy",                     #            regressor
     "einops",                    #            regressor
     "numpy",                     #            regressor analytical
-    "pandas",                    #            regressor
     "celmech>=1.5.0,<1.6.0"      #                      analytical
 ]
 
 [project.optional-dependencies]
 test = [
-    "IPython"
+    "IPython",
+    "matplotlib",
+    "pandas",
+    "jupyter"
 ]
 
 [project.urls]
 Homepage = "https://github.com/dtamayo/spock"
 
+[tool]
+rye = { dev-dependencies = [
+    "ipython>=8.24.0",
+    "jupyter>=1.0.0",
+    "ipykernel>=6.29.4",
+    "matplotlib>=3.9.0",
+] }
+
 [tool.setuptools_scm]
 version_scheme = "release-branch-semver"
 
 [tool.setuptools]
 packages = ["spock"]
-package-data = {spock = ["models/featureclassifier.json", "models/regression/steps=300000*.pkl"]}
+package-data = {spock = ["models/featureclassifier.json", "models/regression/ensemble_part_*"]}

spock/additional_feature_functions.py

@@ -1,6 +1,4 @@
-import rebound
 import numpy as np
-import pandas as pd
 from collections import OrderedDict
 from feature_functions import get_pairs, find_strongest_MMR, populate_trio
 from AMD_functions import AMD, AMD_crit

spock/deepregressor.py

@@ -1,30 +1,19 @@
 import numpy as np
 import math
-from scipy.stats import truncnorm
 import os
 from collections import OrderedDict
 from .tseries_feature_functions import get_extended_tseries
-from copy import deepcopy as copy
 import torch
-from torch import nn
-from torch.nn import Parameter
-from torch.autograd import Variable
-from torch.functional import F
 import glob
-from .spock_reg_model import load_swag
-from pytorch_lightning import Trainer
+from .spock_reg_model import load_swag_safetensors
 import torch
-import time
-import pickle as pkl
 import warnings
 import einops as E
 from scipy.integrate import quad
 from scipy.interpolate import interp1d
-import pytorch_lightning as pl
 from .simsetup import init_sim_parameters
 from multiprocessing import cpu_count
 from multiprocessing.pool import ThreadPool as Pool
-import rebound
 import random
 warnings.filterwarnings('ignore', "DeprecationWarning: Using or importing the ABCs")
 
@@ -127,16 +116,19 @@ def fast_truncnorm(
     return t_inst_samples.reshape(*oldscale.shape)
 
 class DeepRegressor(object):
-    def __init__(self, cuda=False, filebase='*v50_*output.pkl'):
-        super(DeepRegressor, self).__init__()
+    def __init__(self, cuda=False, filebase='ensemble_part_*json'):
+        super().__init__()
         pwd = os.path.dirname(__file__)
         pwd = pwd + '/models/regression'
         self.cuda = cuda
 
         #Load model
+        all_model_param_filenames = glob.glob(pwd + '/' + filebase)
+        all_model_param_filenames.sort()
+        model_basenames = [".".join(fname.split(".")[:-1]) for fname in all_model_param_filenames]
         self.swag_ensemble = [
-            load_swag(fname).cpu()
-            for i, fname in enumerate(glob.glob(pwd + '/' + filebase)) #0.78, 0.970
+            load_swag_safetensors(base_fname)
+            for base_fname in model_basenames
         ]
         #Data scaling parameters:
         self.scale_ = np.array([2.88976974e+03, 6.10019661e-02, 4.03849732e-02, 4.81638693e+01,

spock/modelfitting.py

@@ -1,4 +1,3 @@
-import pandas as pd
 from sklearn.metrics import precision_recall_curve
 from sklearn.metrics import roc_curve, confusion_matrix, auc
 from sklearn import metrics
@@ -12,6 +11,7 @@ def hasnull(row):
         return 1
 
 def train_test_split(trainingdatafolder, features=None, labelname='Stable', filter=False, filtertimes=False):
+    import pandas as pd
     dataset = pd.read_csv(trainingdatafolder+"trainingdata.csv", index_col = 0)
     if features is None:
         features = dataset.columns.values

spock/models/regression/ensemble_part_0.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 0, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_1.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 1, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_10.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 10, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_11.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 11, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_12.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 12, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_13.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 13, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_14.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 14, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_15.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 15, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_16.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 16, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_17.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 17, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_18.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 18, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_19.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 19, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_2.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 2, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_20.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 20, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_21.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 21, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}

spock/models/regression/ensemble_part_22.json

@@ -0,0 +1 @@
+{"hparams": {"seed": 22, "batch_size": 2000, "hidden": 40, "in": 1, "latent": 20, "lr": 0.0005, "swa_lr": 0.0001, "out": 1, "samp": 5, "swa_start": 25000, "weight_decay": 1e-14, "to_samp": 1, "epochs": 1251, "scheduler": true, "scheduler_choice": "swa", "steps": 50000, "beta_in": 1e-05, "beta_out": 0.001, "act": "softplus", "noisy_val": false, "gradient_clip": 0.1, "fix_megno": false, "fix_megno2": true, "include_angles": true, "include_mmr": false, "include_nan": false, "include_eplusminus": false, "power_transform": false, "lower_std": false, "train_all": false, "include_derivatives": false, "time_series_features": 41, "save_freq": 25, "eval_freq": 5, "momentum": 0.9, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30}, "swa_params": {"swa_lr": 0.0001, "swa_start": 25000, "swa_recording_lr_factor": 0.5, "c": 5, "K": 30, "steps": 50000}}