ACID-code 2.0.0a2__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

ACID_code/acid.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import traceback
 import warnings
 warnings.filterwarnings("ignore")
 import sys, emcee, os, time, inspect, inspect, contextlib
@@ -15,6 +16,10 @@ from .result import Result
 from .data import Data, Config, MaskingLines, LineList, DataList
 from .errors import ContinuumError
 from .utils import IntLike, Scalar, Array1D, Array2D
+try:
+    import dynesty
+except ImportError:
+    dynesty = None
 
 @beartype
 class Acid:
@@ -55,9 +60,9 @@ class Acid:
 
         Important note: All defaults in the signature are None, meaning if any values are input, they will override the default :py:class:`Config` and/or :py:class:`Data` values or
         any values that have already been input. The defaults within the config are written below. The config defaults can also be accessed via 
-        :py:attr:`ACID_code.Config.defaults` (returning a dictionary of defaults for both initialisation and run_acid).
+        :py:attr:`ACID_code.Config.defaults` (returning a dictionary of defaults for both initialisation and the ACID method).
 
-        All parameters below and in run_ACID are stored in the :py:class:`Config` instance, unless explicitly stated to be in the :py:class:`Data` instance.
+        All parameters below and in the ACID method are stored in the :py:class:`Config` instance, unless explicitly stated to be in the :py:class:`Data` instance.
         The :py:class:`Config` instance is for runtime settings and the :py:class:`Data` instance is for storing data and any calculations. 
 
         Parameters
@@ -68,12 +73,11 @@ class Acid:
             choose your own velocity grid, by default None, stored in the Data instance.
         linelist : :py:type:`Array2D | str` | :py:class:`LineList` | dict`, optional
             The linelist to use for LSD. The linelist should have wavelengths in angstroms and relative depths between 0 and 1.
-            This is a required parameter if linelist_wl and linelist_depths are not provided. It can be of the forms:
+            This is a required parameter. It can be of the forms:
             - String: A path to a VALD linelist in string format. Support for other linelists may be added in the future or on request.
             - :py:type:`Array2D`: A 2D array-like object indexed such that 0 is wavelengths and 1 is depths.
             - dict: A dictionary with keys "wavelengths" and "depths", each containing array-like objects for the wavelengths and depths respectively.
             - :py:class:`LineList`: The :py:class:`LineList` class is used to expose the linelist for masking or getting/plotting the linelist. You can input an instance if you have one.
-            - If None, linelist_wl and linelist_depths must be provided (see below), by default None, stored in the Data instance.
         order : :py:type:`IntLike`, optional
             If this ACID instance is intended as a run on a specific order, then you can designate this instance for that order. This will allow
             the resulting Data instance to track of which order the profiles correspond to. Note that orders can be indexed by the correct indexing
@@ -96,7 +100,7 @@ class Acid:
             By default 2 (medium).
         sampler_progress : :py:type:`bool`, optional
             A verbosity override for just the MCMC sampling progress.
-            By default None which does not override, but if True/False, it will overwrite with that value.
+            By default None which does not override, but if True/False, it will overwrite with that value, and use/don't use a tqdm output for the sampler.
         masking_lines : :py:type:`dict` | :py:class:`MaskingLines`, optional
             Telluric lines (in angstroms) and widths in (km/s) to mask from the wavelength regions from. Unless you'd like to change the default masking
             lines, we recommend just using the defaults (leaving this as None), which are based on telluric lines and strong hydrogen/metal lines in the 
@@ -112,7 +116,7 @@ class Acid:
             The path to save the sampler HDF5 backend file to.
             If None, the sampler is not saved and only stored in memory. By default None.
             Note that if your path points to an existing file, it will be overwritten on Acid initialization.
-            If True, we use the emcee HDF5 backend to store and load the sampler.
+            If existing, we use the emcee HDF5 backend to store and load the sampler.
             Should be a valid file path that ends with ".h5". If the directory containing it does not exist, it will be created.
             Note that if you later try and save the sampler through the data class, it is converted to a HD5 backend.
         data : :py:class:`Data` | :py:class:`DataList`, optional
@@ -231,6 +235,8 @@ class Acid:
         dev_perc              : IntLike|None                = None,   # Config
         n_sig                 : IntLike|None                = None,   # Config
         skips                 : IntLike|None                = None,   # Config
+        od                    : bool|None                   = None,   # Config
+        sampler_type          : str|None                    = None,   # Config
         parallel              : bool|None                   = None,   # Config
         cores                 : IntLike|None                = None,   # Config
         nwalkers              : IntLike|None                = None,   # Config, then Data just before MCMC
@@ -308,6 +314,14 @@ class Acid:
         skips : :py:type:`IntLike`, optional
             An option to only run acid on one in every n pixels, where n is the integer argument. This is only useful for
             testing to get a quicker result especially for larger wavelength ranges or datasets, by default 1 (no skipping)
+        od : :py:type:`bool`, optional
+            If True, runs ACID in optical depth, otherwise, the LSD methods and ACID fitting is performed in flux. By default None which defaults to True.
+            Note that the whole point of ACID is to run LSD in OD, we highly recommend leaving this unless you specifically want to compare.
+        sampler_type : :py:type:`str`, optional
+            If you really try to wish to use the dynesty nested sampler, you can set this to "dynesty". It is almost entirely unsupported
+            by the rest of the code other than to just get a finished result object, and much slower. We highly recommend using None or "emcee" (default).
+            The only reason I added this was to get the Bayesian evidence for model comparison.
+            If "dynesty" is chosen, the dynesty package needs to be installed, and the nsteps parameter is treated as "nlive" to be passed to the NestedSampler.
         parallel : :py:type:`bool`, optional
             If True uses multiprocessing to calculate the profiles for each frame in parallel, see
             https://acid-code.readthedocs.io/en/stable/using_ACID.html#multiprocessing for more details. By default True
@@ -411,6 +425,8 @@ class Acid:
             "dev_perc"              : dev_perc,
             "n_sig"                 : n_sig,
             "skips"                 : skips,
+            "od"                    : od,
+            "sampler_type"          : sampler_type,
             "parallel"              : parallel,
             "cores"                 : cores,
             "nwalkers"              : nwalkers,
@@ -435,6 +451,14 @@ class Acid:
                 print("Parallel MCMC on Windows is not currently supported. Running MCMC serially.")
             self.config.parallel = False
 
+        if self.config.sampler_type == "dynesty":
+            if dynesty is None:
+                raise ImportError("The 'dynesty' sampler requires the 'dynesty' package to be installed.\nPlease install it with 'pip install dynesty' or choose a different sampler type.")
+        if self.config.sampler_type == "dynesty" and not self.config.deterministic_profile:
+            raise ValueError("The 'dynesty' sampler can only be run with deterministic_profile=True (otherwise you'll be waiting hours for a single result)")
+        if self.config.sampler_type == "dynesty" and self.config.max_steps is not None:
+            raise ValueError("Cannot use max_steps as dynesty already natively supports this with live points, set nsteps=nlive. See the dynesty docs for more details.")
+
         # --- Start of the ACID method ---
 
         # Setup and data validation done in data class and applies skips
@@ -484,8 +508,9 @@ class Acid:
         # The code for telluric masking is contained without the MaskingLines class, which both telluric_lines
         # and hydrogen_lines are instances of.
         line_mask = self.config.masking_lines.get_masks(self.data.wavelengths["combined"])
-        line_mask = np.all(line_mask, axis=0)
-        self.data.errors["combined"][line_mask] = 1e12
+        if line_mask != []:
+            line_mask = np.all(line_mask, axis=0)
+            self.data.errors["combined"][line_mask] = 1e12
 
         # Get the initial polynomial coefficents
         if not hasattr(self.data.wavelengths, "combined_normalized"):
@@ -558,22 +583,25 @@ class Acid:
         self.data.nwalkers = self.data.ndim * 3 if self.config.nwalkers is None else self.config.nwalkers
         rng = np.random.default_rng(self.config.seed)
 
-        # Starting values of walkers with independent variation
-        sigma = 0.8 * 0.005
-        initial_state = []
-        for i in range(0, len(self.data.model_inputs)):
-            if i < len(self.data.velocities):
-                if not self.config.deterministic_profile:
-                    pos = rng.normal(self.data.model_inputs[i], sigma, (self.data.nwalkers, ))
+        # Starting values of walkers with independent variation#
+        if self.config.sampler_type == "emcee":
+            sigma = 0.8 * 0.005
+            initial_state = []
+            for i in range(0, len(self.data.model_inputs)):
+                if i < len(self.data.velocities):
+                    if not self.config.deterministic_profile:
+                        pos = rng.normal(self.data.model_inputs[i], sigma, (self.data.nwalkers, ))
+                    else:
+                        continue
                 else:
-                    continue
-            else:
-                x1 = self.data.model_inputs[i]
-                rounded_sigma = round(x1, 1-int(floor(log10(abs(x1))))-1)
-                sigma = abs(rounded_sigma) / 10
-                pos = rng.normal(self.data.model_inputs[i], sigma, (self.data.nwalkers, ))
-            initial_state.append(pos)
-        initial_state = np.array(initial_state).T
+                    x1 = self.data.model_inputs[i]
+                    rounded_sigma = round(x1, 1-int(floor(log10(abs(x1))))-1)
+                    sigma = abs(rounded_sigma) / 10
+                    pos = rng.normal(self.data.model_inputs[i], sigma, (self.data.nwalkers, ))
+                initial_state.append(pos)
+            initial_state = np.array(initial_state).T
+        else:
+            initial_state = None
 
         ### ACID initialialised ###
         self.data.setup_time += time.time() - init_t0
@@ -617,12 +645,12 @@ class Acid:
         """
         This method is no longer supported in ACID. Please use the ACID function with the appropriate inputs for HARPS spectra instead. 
         Future versions of ACID will provide functions to load and configure data from a range of different standard instruments. 
-        If you still really wish to use ACID_HARPS, the last stable version of ACID with the method is 1.4.5. Try: pip install ACID_code==1.4.5
+        If you still really wish to use ACID_HARPS, the last stable version of ACID with the method is 1.4.5. Try: pip install ACID_code_v2==1.4.5
         """
         raise NotImplementedError(f"ACID_HARPS is no longer supported in ACID. \n"
         f"Please use the ACID function with the appropriate inputs for HARPS spectra instead. \n"
         f"Future versions of ACID will provide functions to load and configure data from a range of different standard instruments. \n"
-        f"If you still really wish to use ACID_HARPS, the last stable version of ACID with the method is 1.4.5. Try: pip install ACID_code==1.4.5")
+        f"If you still really wish to use ACID_HARPS, the last stable version of ACID with the method is 1.4.5. Try: pip install ACID_code_v2==1.4.5")
 
     def combine_spec(
         self,
@@ -834,9 +862,12 @@ class Acid:
             self.data.plot_continuum_fit(plot_type=plot_type)
 
         if np.any(flux_obs <= 0) or np.any(new_errors <= 0):
-            raise ContinuumError("Continuum fit resulted in non-positive flux or errors, which is not physical.\n " \
+            error = ContinuumError("Continuum fit resulted in non-positive flux or errors, which is not physical.\n " \
             "Consider adjusting the polynomial order or continuum percentile. Use verbose=3 to see the plot of the continuum fit.\n " \
             "Note that this will only work for interactive terminals or displays which work with plt.show()")
+            self.data.exception = error
+            self.data.traceback = traceback.format_stack()
+            raise error
 
         return poly_coeffs, flux_obs, new_errors
 
@@ -856,7 +887,7 @@ class Acid:
         sn = self.data.sn["combined"]
 
         # Use the initial LSD run to get the forward model and scaled residuals
-        forward, _profile = mcmc.MCMC(x, y, yerr, self.data.alpha).full_model(self.data.model_inputs)
+        forward, _profile = mcmc.MCMC(x, y, yerr, self.data.alpha, od=self.config.od).full_model(self.data.model_inputs)
         residuals = (y - forward) / forward
 
         # Chunk masking based on deviation from residuals
@@ -953,7 +984,8 @@ class Acid:
         """
 
         # Get default sampler kwargs from initial state
-        sampler_kwargs, mcmc_kwargs = self._get_sampler_kwargs(nsteps, state)
+        if self.config.sampler_type == "emcee":
+            sampler_kwargs, mcmc_kwargs = self._get_sampler_kwargs(nsteps, state)
         pool_context = nullcontext(None)
 
         if self.config.parallel:
@@ -964,14 +996,25 @@ class Acid:
             
             ctx = mp.get_context("fork")
             pool_context = ctx.Pool(processes=self.config.cores, initializer=mcmc._mp_init_worker, initargs=(self.data,))
-            log_prob_fn = mcmc._mp_log_probability
+            log_prob = mcmc._mp_log_probability if self.config.sampler_type == "emcee" else mcmc._mp_log_likelihood
+            ptform = mcmc._mp_ptform
+            queue_size = os.cpu_count()
         else:
             MCMC = mcmc.MCMC(self.data)
-            log_prob_fn = MCMC
-        
+            log_prob = MCMC if self.config.sampler_type == "emcee" else MCMC.dynesty_logprob
+            ptform = MCMC.ptform
+            queue_size = None
+
         with pool_context as pool:
-            self.sampler = EnsembleSampler(log_prob_fn=log_prob_fn, pool=pool, **sampler_kwargs)
-            self.sampler.run_mcmc(**mcmc_kwargs)
+            if self.config.sampler_type == "emcee":
+                self.sampler = EnsembleSampler(log_prob_fn=log_prob, pool=pool, **sampler_kwargs)
+                self.sampler.run_mcmc(**mcmc_kwargs)
+            else:
+                import dynesty
+                if self.config.parallel:
+                    pool.size = self.config.cores
+                self.sampler = dynesty.NestedSampler(log_prob, ptform, self.data.ndim, self.config.nsteps, pool=pool, queue_size=queue_size)
+                self.sampler.run_nested(print_progress=self.config.verbose>1)
 
     def run_mcmc_until_converged(self, max_steps:IntLike, state=None) -> None:
         """