Cleanup

nanograv · Feb 2, 2024 · 5c06e3e · 5c06e3e
1 parent 4704d4e
commit 5c06e3e
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 28 deletions.
diff --git a/holodeck/__init__.py b/holodeck/__init__.py
@@ -34,7 +34,7 @@
 """
 
 __author__ = "NANOGrav"
-__copyright__ = "Copyright (c) 2023 NANOGrav"
+__copyright__ = "Copyright (c) 2024 NANOGrav"
 __license__ = "MIT"
 
 import os

diff --git a/holodeck/librarian/posterior_populations.py b/holodeck/librarian/posterior_populations.py
@@ -16,7 +16,15 @@
                                    |     |-------------------> 30 frequency bins
                                    |-------------------------> 20 years observing baseline = 1/(20yr) lowest frequency
 
+To-Do
+-----
+* Improve handling of data path.
+* Improve handling/specification of parameter space.
+    * Allow changes to be passed in through API and or CL
+    * Make each particular 15yr dataset specify its own parameter space (these need to match up anyway!)
+
 """
+
 import argparse
 from pathlib import Path
 import numpy as np
@@ -33,6 +41,8 @@
 NLOUDEST = 10
 
 # Path to chains, fitting holodeck populations to data, giving parameter posteriors
+# This is the `15yr_astro_data` currently stored on google drive:
+# https://drive.google.com/drive/u/1/folders/1wFy_go_l8pznO9D-a2i2wFHe06xuOe5B
 PATH_DATA = Path(
     "/Users/lzkelley/Programs/nanograv/15yr_astro_data/"
     "phenom/ceffyl_chains/astroprior_hdall/"
@@ -42,8 +52,7 @@
 PSPACE = holo.librarian.param_spaces_classic.PS_Classic_Phenom_Uniform
 
 # Path to save output data
-PATH_OUTPUT = Path(__file__).parent.joinpath("output")
-assert PATH_OUTPUT.is_dir(), f"{PATH_OUTPUT=} does not exist!"
+PATH_OUTPUT = Path(holo._PATH_OUTPUT).resolve().joinpath("15yr_pops")
 
 
 def main(args=None):
@@ -57,15 +66,17 @@ def main(args=None):
         args = setup_argparse()
 
     # load chains (i.e. parameter posterior distributions)
-    chains = load_chains(PATH_DATA)
+    # chains = load_chains(PATH_DATA)
 
     # select parameters for this population
     if args.maxlike:
         pkey = "ML"
-        pars = get_maxlike_pars_from_chains(chains)
+        # pars = get_maxlike_pars_from_chains(chains)
+        pars = get_maxlike_pars_from_chains()
     else:
         pkey = "draw"
-        pars = sample_chains(chains)
+        # pars = sample_pars_from_chains(chains)
+        pars = sample_pars_from_chains()
 
     # construct output filename
     output = Path(args.output).resolve()
@@ -78,15 +89,17 @@ def main(args=None):
         if not fname.exists():
             break
         if (num > 0) and args.maxlike and (ml_warning is False):
-            holo.log.warning("")
+            err = "Maximum likelihood population with these paramters already exists!  {fname}"
+            holo.log.error(err)
+            raise RuntimeError(err)
 
     else:
         raise RuntimeError(f"Could not find a filename that doesn't exist!  e.g. {fname}")
 
     # ---- Construct population and derived properties
 
     # Build populations with holodeck
-    data = load_population_for_pars(args, pars)
+    data, classes = load_population_for_pars(args, pars)
 
     # ---- Save to output file
 
@@ -165,7 +178,7 @@ def load_population_for_pars(args, pars):
         Typically `args` should be loaded using the `setup_argparse` function.
     pars : dict
         Binary population parameters for the appropriate parameter space `PSPACE`.
-        Typically the `pars` should be loaded using either the `sample_chains` or the
+        Typically the `pars` should be loaded using either the `sample_pars_from_chains` or the
         `get_maxlike_pars_from_chains` function.
 
     Returns
@@ -231,7 +244,12 @@ def load_population_for_pars(args, pars):
         mtot_edges=edges[0], mrat_edges=edges[1], redz_edges=edges[2], fobs_orb_edges=edges[3],
     )
 
-    return data
+    classes = dict(
+        sam=sam,
+        hard=hard,
+    )
+
+    return data, classes
 
 
 def load_chains(path_data):
@@ -272,7 +290,7 @@ def load_chains(path_data):
     return data
 
 
-def sample_chains(chains):
+def sample_pars_from_chains(chains=None):
     """Sample randomly from the given chains (i.e. parameter posteriors).
 
     Arguments
@@ -291,13 +309,16 @@ def sample_chains(chains):
             'mmb_mamp_log10', 'mmb_scatter_dex', 'hard_gamma_inner'],
 
     """
+    if chains is None:
+        chains = load_chains(PATH_DATA)
+
     nlinks = list(chains.values())[0].size
     idx = np.random.choice(nlinks)
     pars = {key: value[idx] for key, value in chains.items()}
     return pars
 
 
-def get_maxlike_pars_from_chains(chains):
+def get_maxlike_pars_from_chains(chains=None):
     """Load the maximum-likelihood (ML) parameters from the given chains (i.e. parameter posteriors).
 
     KDEs from `kalepy` are used to construct the ML parameters.
@@ -319,6 +340,8 @@ def get_maxlike_pars_from_chains(chains):
 
     """
     import kalepy as kale
+    if chains is None:
+        chains = load_chains(PATH_DATA)
 
     # Get maximum likelihood parameters (estimate using KDE)
     mlpars = {}

diff --git a/notebooks/nanograv-15yr-populations.ipynb b/notebooks/nanograv-15yr-populations.ipynb
@@ -24,17 +24,24 @@
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
     "import holodeck as holo\n",
+    "from holodeck.librarian import posterior_populations\n",
     "from holodeck.constants import MSOL, YR"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "import sys\n",
-    "print(sys.argv)"
+    "# Load data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### from save file\n",
+    "\n",
+    "(generated by `holodeck/librarian/posterior_populations.py`)"
    ]
   },
   {
@@ -43,18 +50,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import gen_holodeck_pops\n",
-    "args = gen_holodeck_pops.setup_argparse(\"\")\n",
-    "print(args)\n",
-    "args = gen_holodeck_pops.setup_argparse(\"\", nloudest=12, nreals=6, maxlike=True)\n",
-    "print(args)"
+    "# Specify filename to load (produced from `gen_holodeck_pops.py` script)\n",
+    "fname = \"./output/t20.0yr_nf30_nr100_nl10_ML_0000.npz\"\n",
+    "# Load data\n",
+    "data = np.load(fname)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Load data"
+    "### generate from posterior chains"
    ]
   },
   {
@@ -63,10 +69,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Specify filename to load (produced from `gen_holodeck_pops.py` script)\n",
-    "fname = \"./output/t2.0yr_nf12_nr7_nl6_draw.npz\"\n",
-    "# Load data\n",
-    "data = np.load(fname)"
+    "args = posterior_populations.setup_argparse(\"\", maxlike=True)\n",
+    "pars = posterior_populations.get_maxlike_pars_from_chains()\n",
+    "data, classes = posterior_populations.load_population_for_pars(args, pars)"
    ]
   },
   {
@@ -259,7 +264,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.11.7"
   }
  },
  "nbformat": 4,