new toy example, review fixes, sharding

homerjed · Jan 4, 2025 · 2bdc33b · 2bdc33b
1 parent f5eec6a
commit 2bdc33b
Show file tree

Hide file tree

Showing 12 changed files with 821 additions and 130 deletions.
diff --git a/.gitignore b/.gitignore
@@ -16,4 +16,4 @@ __pycache__/
 .DS_Store
 
 affine.py
-_experiment.py
+_experiment.py
diff --git a/README.md b/README.md
@@ -92,7 +92,7 @@ and have a look at [examples](https://github.com/homerjed/sbiax/tree/main/exampl
 
 ### Citation
 
-If you found this library to be useful in academic work, then please cite:  <!--([arXiv link](https://arxiv.org/abs/2111.00254)) -->
+If you found this library to be useful in academic work, please cite:  <!--([arXiv link](https://arxiv.org/abs/2111.00254)) -->
 
 ```bibtex
 @misc{homer2024simulationbasedinferencedodelsonschneidereffect,

diff --git a/data/shear.py b/data/shear.py
@@ -6,10 +6,8 @@
 import tensorflow_probability.substrates.jax.distributions as tfd
 
 
-def get_shear_experiment():
-    data_dir = "/Users/Jed.Homer/phd/sbiax/data/shear/"
-
-    covariance   = np.loadtxt(data_dir + "covariance_cosmic_shear_PMEpaper.dat")
+def get_shear_experiment(data_dir):
+    covariance   = np.loadtxt(data_dir + "covariance_cosmic_shear_PMEpaper.dat") # / (18. / 5.)
     precision    = np.linalg.inv(np.matrix(covariance))
     mu           = np.loadtxt(data_dir + "DES_shear-shear_a1.0_b0.5_data_vector")[:, 1]
     derivatives  = np.loadtxt(data_dir + "derivatives.dat").T
@@ -39,7 +37,6 @@ def get_shear_experiment():
 
 
 def linearized_model(_alpha, mu, alpha, derivatives):
-    """ Linearised model always uses true mu_0, C """
     return mu + jnp.dot(_alpha - alpha, derivatives)
 
 
@@ -100,7 +97,7 @@ def _mle(d, pi, Finv, mu, dmu, precision):
     return pi + jnp.linalg.multi_dot([Finv, dmu, precision, d - mu])
 
 
-def get_experiment_data(key, true_covariance, n_sims, *, results_dir):
+def get_experiment_data(key, true_covariance, n_sims, *, results_dir, data_dir):
 
     key, key_prior, key_simulate = jr.split(key, 3)
 
@@ -116,7 +113,7 @@ def get_experiment_data(key, true_covariance, n_sims, *, results_dir):
         Finv, 
         lower,
         upper
-    ) = get_shear_experiment()
+    ) = get_shear_experiment(data_dir=data_dir)
 
     # Estimate covariance, Fisher information and precision given n_sims 
     if true_covariance:

diff --git a/examples/shear.ipynb b/examples/shear.ipynb
diff --git a/examples/shear_nn.ipynb b/examples/shear_nn.ipynb
diff --git a/examples/shear_optuna.ipynb b/examples/shear_optuna.ipynb