Added prior selection trick for speed

README.rst

@@ -22,6 +22,7 @@ PBjam is toolbox for analyzing the oscillation spectra of solar-like oscillators
 
 The mode identification works by fitting the asymptotic relation for p-modes to the l=2,0 pairs, which is followed by a applying selection of models for fitting the l=1 modes where each model is suitable for different stages of evolution.
 The process relies on of large set of previous observations of the model parameters, which are then used to construct a prior distribution to inform the sampling. The observations have been gathered from the Kepler, K2 and TESS missions, and expanding it to improve accuracy is an on-going process. 
+The ``runl20model`` and ``runl1model`` mode-identification calls accept a ``loglikelihoodMultiplier`` argument, defaulting to ``1.0``, which scales the final log-likelihood value passed to the nested sampler.
 
 Modeling the modes, or 'peakbagging', is done using the a nested sampling or MCMC algorithm, where Lorentzian profiles are fit to each of the identified modes, with much fewer constraints than during the mode ID process. This allows for a more accurate model of the spectrum of frequencies than the heavily parameterized models like the asymptotic relations.
 

docs/source/MSmodels.rst

@@ -0,0 +1,4 @@
+MSmodels
+^^^^^^^^
+.. automodule:: pbjam.MSmodels
+    :members:

docs/source/api.rst

@@ -18,6 +18,7 @@ Modules in PBjam
    jar
    l1models
    l20models
+   MSmodels
    modeID
    peakbagging
    plotting
@@ -31,4 +32,3 @@ Modules in PBjam
 
 
 
-

docs/source/conf.py

@@ -15,6 +15,12 @@
 import os
 import sys
 sys.path.insert(0, os.path.abspath('../..'))
+from importlib.metadata import PackageNotFoundError, version as package_version
+
+try:
+    pbjam_version = package_version('pbjam')
+except PackageNotFoundError:
+    pbjam_version = '0+unknown'
 
 
 # -- Project information -----------------------------------------------------
@@ -24,9 +30,9 @@
 author = 'Martin Nielsen'
 
 # The short X.Y version
-version = ''
+version = '.'.join(pbjam_version.split('.')[:2])
 # The full version, including alpha/beta/rc tags
-release = '0.0.1'
+release = pbjam_version
 
 
 # -- General configuration ---------------------------------------------------
@@ -52,12 +58,11 @@
     'nbsphinx',
 ]
 
-
-
-# If your notebooks contain non-trivial outputs like images, set this option
-nbsphinx_allow_errors = True  # This ensures Sphinx builds even if a notebook fails
+# Keep notebook conversion failures visible during docs builds.
+nbsphinx_allow_errors = False
 
-# Run the notebooks or never.
+# Tutorials download mission data and run expensive samplers, so docs builds
+# render the checked-in outputs instead of re-executing notebooks.
 nbsphinx_execute = 'never'
 
 # Exclude build directory and Jupyter backup files:
@@ -69,9 +74,10 @@
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
 #
-# source_suffix = ['.rst', '.md']
-#source_suffix = '.rst'
-source_suffix = {'.rst': 'restructuredtext'}
+source_suffix = {
+    '.rst': 'restructuredtext',
+    '.ipynb': 'nbsphinx',
+}
 
 # The master toctree document.
 master_doc = 'index'
@@ -83,11 +89,6 @@
 # Usually you set "language" from the command line for these cases.
 language = 'en'
 
-# List of patterns, relative to source directory, that match files and
-# directories to ignore when looking for source files.
-# This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = []
-
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = None
 
@@ -208,4 +209,4 @@
 # If true, `todo` and `todoList` produce output, else they produce nothing.
 todo_include_todos = True
 
-autodoc_default_flags = ['members', 'special-members']
+autodoc_default_flags = ['members', 'special-members']

docs/source/examples.rst

@@ -0,0 +1,16 @@
+Examples
+========
+
+These notebooks demonstrate the main PBjam workflows.
+
+.. toctree::
+   :maxdepth: 1
+
+   Examples/example-session.ipynb
+   Examples/example-star.ipynb
+   Examples/example-modeID.ipynb
+   Examples/example-peakbag.ipynb
+   Examples/example-plotting.ipynb
+   Examples/example-docstrings
+   Examples/example-unittests
+   Examples/example-githubworkflow

docs/source/index.rst

@@ -6,6 +6,7 @@
 
    setup
    usage
+   examples
    api
    license
    help

docs/source/usage.rst

@@ -5,7 +5,7 @@ PBjam is intended to be a user-friendly peakbagging tool. The most straightforwa
 
 Session
 -------
-The :class:`~pbjam.core.session` class is the most straightforward way to analyze one or more stars with PBjam. It can automatically download the data and compute the power density spectrum, and then go through all the steps in the moden ID and peakbagging process. The `Session notebook <Examples/example-session.ipynb>`_ provides an example of how to use the :class:`~pbjam.core.session` class. 
+The :class:`~pbjam.core.session` class is the most straightforward way to analyze one or more stars with PBjam. It can automatically download the data and compute the power density spectrum, and then go through all the steps in the mode ID and peakbagging process. The `Session notebook <Examples/example-session.ipynb>`_ provides an example of how to use the :class:`~pbjam.core.session` class. 
 
 
 Star
@@ -20,8 +20,8 @@ Advanced
 It is not strictly necessary to use either the :class:`~pbjam.core.session` or :class:`~pbjam.core.star` classes. The `mode ID <Examples/example-modeID.ipynb>`_ and `peakbag <Examples/example-peakbag.ipynb>`_ notebooks show a lower-level walkthrough of the steps that PBjam goes through for peakbagging.
 
 .. note:: 
-    For additional useful examples see the `Examples <https://github.com/grd349/PBjam/tree/master/Examples>`_ directory.
+    For additional useful examples see the `Examples <examples.html>`_ page.
 
 Papers
 ------
-We have published a few papers on various bits. The `first paper <https://ui.adsabs.harvard.edu/abs/2021AJ....161...62N/abstract>`_ provides information mainly on the initial version of PBjam. The `second paper <https://ui.adsabs.harvard.edu/abs/2023A%26A...676A.117N/abstract>`_ discusses the method used to construct the prior probability densities that we now use in the latest version of PBjam. The latest paper (in prep.) focuses on how we construct the models for the l=1 modes.
+We have published a few papers on various bits. The `first paper <https://ui.adsabs.harvard.edu/abs/2021AJ....161...62N/abstract>`_ provides information mainly on the initial version of PBjam. The `second paper <https://ui.adsabs.harvard.edu/abs/2023A%26A...676A.117N/abstract>`_ discusses the method used to construct the prior probability densities that we now use in the latest version of PBjam. The latest paper (in prep.) focuses on how we construct the models for the l=1 modes.

pbjam/DR.py

@@ -309,7 +309,7 @@ def fit_weightedPCA(self, dim):
 
         self.covariance = self.covarianceMatrix(_X)
 
-        self.eigvals, self.eigvectors = jnp.linalg.eig(self.covariance)
+        self.eigvals, self.eigvectors = jnp.linalg.eigh(self.covariance)
 
         self.sortidx = sorted(range(len(self.eigvals)), key=lambda i: self.eigvals[i], reverse=True)[:self.dimsR]
 
@@ -319,6 +319,29 @@ def fit_weightedPCA(self, dim):
 
         self.dataR = self.transform(self.dataF)
 
+    def setLatentNormalPrior(self, latentSample=None):
+        """Set independent normal priors for each latent PCA coordinate."""
+
+        from pbjam import distributions as dist
+
+        if latentSample is None:
+            latentSample = self.dataR
+
+        latentSample = np.asarray(latentSample)
+        loc = np.mean(latentSample, axis=0)
+        scale = np.std(latentSample, axis=0)
+        scale = np.where(np.isfinite(scale) & (scale > 0), scale, 1.0)
+
+        self.latentPriorLoc = jnp.array(loc)
+        self.latentPriorScale = jnp.array(scale)
+        self.latentPriors = [dist.normal(loc=self.latentPriorLoc[i],
+                                         scale=self.latentPriorScale[i])
+                             for i in range(self.dimsR)]
+        self.ppf = [prior.ppf for prior in self.latentPriors]
+        self.pdf = [prior.pdf for prior in self.latentPriors]
+        self.logpdf = [prior.logpdf for prior in self.latentPriors]
+        self.cdf = [prior.cdf for prior in self.latentPriors]
+
     def covarianceMatrix(self, _X):
         """ Compute the weighted covariance matrix
 
@@ -338,3 +361,89 @@ def covarianceMatrix(self, _X):
         C = _X.T@W@_X * jnp.sum(self.weights) / (jnp.sum(self.weights)**2 - jnp.sum(self.weights**2))
 
         return C
+
+    def refinePriorByObservables(self, N=10000, minAccepted=100, sigmaInflation=1,
+                                 rng=None):
+        """Refit latent normal priors using draws consistent with observations.
+
+        Draws from the current latent prior, maps those points into the PCA
+        parameter space, accepts them with a normal likelihood in the observed
+        dimensions, and rebuilds the independent latent normal priors from the
+        accepted points.
+        """
+
+        if not hasattr(self, 'ppf'):
+            raise AttributeError('Set the initial latent prior before refining the prior.')
+
+        if not hasattr(self, 'dimsR') or self.dimsR == 0:
+            return None
+
+        obsLabels = [key for key in self.selectLabels
+                     if (key in self.varLabels) and (key in self.obs)]
+
+        obsLabels = [key for key in obsLabels if self.obs[key][1] > 0]
+
+        if len(obsLabels) == 0:
+            warnings.warn('Selective prior refinement skipped: no observed selection labels are in the PCA variables.',
+                          stacklevel=2)
+            return None
+
+        if rng is None:
+            rng = np.random.default_rng()
+        elif not hasattr(rng, 'uniform'):
+            rng = np.random.default_rng(rng)
+
+        obsIdx = np.array([self.varLabels.index(key) for key in obsLabels])
+        obsVals = np.array([self.obs[key][0] for key in obsLabels])
+        baseObsErrs = np.array([self.obs[key][1] for key in obsLabels])
+
+        nDraws = max(1, int(N))
+        currentSigmaInflation = sigmaInflation
+        while True:
+            latentDraws = rng.normal(loc=np.asarray(self.latentPriorLoc),
+                                     scale=np.asarray(self.latentPriorScale),
+                                     size=(nDraws, self.dimsR))
+
+            physicalDraws = np.asarray(self.inverse_transform(jnp.array(latentDraws)))
+
+            finite = np.all(np.isfinite(physicalDraws), axis=1)
+            obsErrs = baseObsErrs * currentSigmaInflation
+            delta = (physicalDraws[:, obsIdx] - obsVals) / obsErrs
+            logLike = -0.5 * np.sum(delta**2, axis=1)
+            finite &= np.isfinite(logLike)
+
+            if not np.any(finite):
+                raise ValueError('Selective prior refinement found no finite prior draws.')
+
+            acceptProb = np.zeros_like(logLike)
+            acceptProb[finite] = np.exp(logLike[finite] - np.max(logLike[finite]))
+            accepted = rng.uniform(0, 1, size=len(logLike)) < acceptProb
+
+            M = int(np.sum(accepted))
+            self.selectivePriorInfo = {'draws': nDraws,
+                                       'accepted': M,
+                                       'minAccepted': int(minAccepted),
+                                       'sigmaInflation': currentSigmaInflation,
+                                       'labels': obsLabels}
+
+            if M >= minAccepted:
+                break
+
+            nextNDraws = 2 * nDraws
+            nextSigmaInflation = 2 * currentSigmaInflation
+            warnings.warn(f'Selective prior refinement accepted {M} points, fewer than minAccepted={minAccepted}. '
+                          f'Retrying with {nextNDraws} draws and sigmaInflation={nextSigmaInflation}.',
+                          stacklevel=2)
+            nDraws = nextNDraws
+            currentSigmaInflation = nextSigmaInflation
+
+        self.selectivePhysicalSample = physicalDraws[accepted, :]
+        self.selectiveSubset = pd.DataFrame(self.selectivePhysicalSample, columns=self.varLabels)
+
+        selectiveLatentSample = np.asarray(self.transform(jnp.array(self.selectivePhysicalSample)))
+        self.selectiveLatentSample = selectiveLatentSample
+
+        self.dataR = jnp.array(selectiveLatentSample)
+        self.setLatentNormalPrior(selectiveLatentSample)
+
+        return selectiveLatentSample

pbjam/IO.py

@@ -301,7 +301,8 @@ def windowfunction(self, df, width=None, oversampling=10):
         if width is None:
             width = 100*df
 
-        freq_cen = 0.5*self.Nyquist
+        # freq_cen = 0.5*self.Nyquist
+        freq_cen = max([0.5*self.Nyquist, width + df])
 
         Nfreq = int(oversampling*width/df)
 
@@ -554,4 +555,3 @@ def _getPriorPath():
 
     return os.path.join(*[PACKAGEDIR, 'data', 'prior_data.csv'])
 
-