ACID-code 2.0.0a2__py3-none-any.whl → 2.0.0a3__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:
@@ -96,7 +101,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 
@@ -231,6 +236,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 +315,11 @@ 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)
+        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 +423,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 +449,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 +506,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 +581,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 +643,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 +860,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 +885,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 +982,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 +994,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:
         """

ACID_code/data.py

@@ -207,6 +207,8 @@ class Config:
         "dev_perc" : 25,
         "n_sig" : 3,
         "skips" : 1,
+        "od"    : True,
+        "sampler_type" : "emcee",
         "parallel" : True,
         "cores" : None,
         "nwalkers" : None,
@@ -518,6 +520,12 @@ class Data:
     combined_profile  : Optional[list] = None
     #: The final fitted continuum model and errors
     continuum_model   : Optional[np.ndarray] = None
+    #: The forward model using the final profile, alpha matrix, and continuum model
+    forward_model     : Optional[np.ndarray] = None
+    #: Errors on the above forward model, usually not needed
+    forward_errors    : Optional[np.ndarray] = None
+    #: The x-axis for the above forward model, which is just the combined wavelength grid, and set in Result.process_results
+    forward_x        : Optional[np.ndarray] = None
     #: The number of steps taken in the MCMC sampling, used for checking convergence and for resuming
     nsteps            : Optional[int]  = 0
     #: A flag for whether the profiles have been fully calculated to avoid recalculating
@@ -535,6 +543,10 @@ class Data:
     results_time       : Optional[float] = 0 
     #: total_time (float) - The total time for the full run
     total_time         : Optional[float] = 0
+    #: The exception class if an error was raised during the run
+    exception          : Optional[Exception] = None
+    #: The traceback string if an error was raised during the run
+    traceback          : Optional[str] = None
 
     # Initialise the properties
     # -------------------------
@@ -573,14 +585,18 @@ class Data:
             if os.path.exists(sampler):
                 self._sampler = utils.backend_to_sampler(HDFBackend(sampler), log_prob_fn)
             else:
-                raise ValueError(f"The provided sampler path '{sampler}' does not exist.")
+                if self.config.verbose > 0:
+                    print(f"Warning: The sampler was not found at the provided path '{sampler}', it may have been moved or deleted. \n"
+                          f"The sampler will be set to None.", flush=True)
+                self._sampler = None
+                # TODO: Allow sampler to have completed results, but no sampler, and configured methods with _requiresampler property that need them
         elif sampler is None:
             if self.config.verbose > 0 and self._sampler is not None:
                 print("Warning, you have discarded the sampler.")
             self._sampler = None
         
         if self._sampler is not None and isinstance(self._sampler.backend, HDFBackend):
-            self.config.sampler_path = self._sampler.backend.filename
+            self.config.sampler_path = os.path.abspath(self._sampler.backend.filename)
 
     @property
     def velocities(self) -> Array1D|None:
@@ -1196,7 +1212,20 @@ class Data:
         """
         Converts the data object to a dictionary payload for saving. This is used internally in the save method, 
         but can also be used for debugging or other purposes.
+
+        Parameters
+        ----------
+        store_sampler : bool, optional
+            Whether to include the MCMC sampler in the dictionary payload, by default True.
+        size_limit : Scalar | None, optional
+            A hard size limit to the sampler in GB.
+            If the sampler exceeds this size, it will not be stored regardless of the store_sampler flag.
+            This is to avoid accidentally storing very large samplers. If None, no limit is set. Default is 1GB.
+            A warning will be printed if this size_limit forces the store_sampler to be False if store_sampler was set to True.
         """
+        if self.sampler is not None and self.config.sampler_type == "dynesty":
+            raise ValueError("Storing the sampler is not currently supported for dynesty samplers.\n" \
+            "If you really want to, separate the sampler with data.sampler.save('sampler') and add it back later.\n")
 
         payload: dict[str, Any] = {}
         for f in fields(self):
@@ -1239,14 +1268,26 @@ class Data:
 
         # Handle sampler separately
         self.sampler = payload.get("sampler", None) # property handles the loading of the sampler
-        
+        if self.sampler is None and self.config.sampler_path is not None:
+            try:
+                self.sampler = self.config.sampler_path
+            except:
+                self.sampler = None
+
         return self
 
     @property
     def result(self):
+        if self.exception is not None:
+            if self.config.verbose > 0:
+                print(f"An exception was raised during the run, cannot return results object.\n"
+                      f"Returning None instead.")
+                return None
         if not self.complete:
-            raise ValueError(f"Results have not yet been calculated, cannot return results object.\n"
-                             f"Please run the MCMC sampling and process the results first.")
+            if self.config.verbose > 0:
+                print(f"Results for order {self.config.order} have not yet been calculated, cannot return results object.\n"
+                      f"Returning None instead.")
+                return None
         from .result import Result
         return Result(self)
 
@@ -1403,8 +1444,9 @@ class DataList:
         order_range      : Array1D|None                   = None,
         config           : Config|list[Config]|None       = None,
         save_dir         : str|None                       = None,
+        overwrite        : bool                           = False,
         verbose          : IntLike|bool|str|None          = None,
-        load                                              = None,
+        _load                                             = None,
         _data_list       : list[Data]|None                = None,
         **config_kwargs,
         ) -> None:
@@ -1453,11 +1495,17 @@ class DataList:
             By default the DataList will save data.pkl and sampler.h5 to the directory (named by the order number) to in this directory.
             If the Configs or kwargs passed contain their own save_path or sampler_path (see :py:class:`Acid`), those instead are used.
             If None, no saving will be done, this is however, not recommended. Default is None.
+        overwrite : bool, optional
+            Whether to overwrite existing with new Data instances when using run_ACID, or to load and use existing Data instance if they exist.
+            If True, if a Data instance already exists for an order, it will be overwritten with the new Data instance generated from the ACID run for that order.
+            Note, that the saving of this new Data instance only applies when run_ACID is run, otherwise it is just held in memory.
+            If False, if a Data instance already exists for an order, it will be loaded and used instead of generating a new Data instance from the ACID run for that order.
+            Default is False.
         verbose : int | bool | str | None, optional
             The verbosity level for printing information during the initialization. 
             Follows the same format as the "verbose" input in the :py:class:`Config` class. 
             Default is None.
-        load : Any, optional
+        _load : Any, optional
             Not yet implemented, do not use. The idea is that you can input a Load object which has its own tools to pull s2d data from common instruments 
             such as ESPRESSO, HARPS, etc. If you want to use this feature, please open an issue or contribute a pull request with the implementation.
         _data_list : list[:py:class:`Data`] | None, optional
@@ -1471,7 +1519,7 @@ class DataList:
         """
 
         # Raise if load was used
-        if load is not None:
+        if _load is not None:
             raise NotImplementedError(f"The 'load' argument is not yet implemented. \n"
                                       f"The idea is that you can input a Load object which has its own tools to pull s2d data from common "\
                                       f"instruments such as ESPRESSO, HARPS, etc. \nIf you want to use this feature, please open an issue or "\
@@ -1485,19 +1533,19 @@ class DataList:
         # Configure verbosity
         self.verbose = Config(verbose=verbose).verbose
 
-        # Configure velocities
+        # All orders should have the same velocity grid and line list
         self.velocities = velocities
 
         # Configure order_range, creates one if not input from the shape of wavelengths
         self.order_range = order_range # if None, will be set later, otherwise self.from_datalist handles the range from configs     
 
-        # Configure save_dir, for saving intermediate results and figures per order
+        # Set class attributes
         self._save_dir = None
-        self.save_dir = save_dir if save_dir is not None else None
-
-        # Set empty class attributes
+        self._data_list = None
         self._combined_profile = None
+        self.overwrite = overwrite
         self.excluded_orders = []
+        self.save_dir = save_dir
 
         if _data_list is not None:
             self.data_list = _data_list # datalist property handles the rest
@@ -1548,11 +1596,23 @@ class DataList:
             data.velocities = velocities
 
             if self.save_dir is not None:
-                save_path = os.path.join(self.save_dir, f"order_{order}", "data.pkl")
-                sampler_path = os.path.join(self.save_dir, f"order_{order}", "sampler.h5")
-                data.config.update_lowpri(save_path=save_path, # set default save path for this order which can be overwritten by user
+                save_path = os.path.abspath(os.path.join(self.save_dir, f"order_{order}", "data.pkl"))
+                sampler_path = os.path.abspath(os.path.join(self.save_dir, f"order_{order}", "sampler.h5"))
+                data.config.update_hipri(save_path=save_path, # set default save path for this order which can be overwritten by user
                                         sampler_path=sampler_path) # set default sampler path for this order which can be overwritten by user
 
+                # Check if file already exists
+                if os.path.exists(data.config.save_path):
+                    if self.overwrite:
+                        if self.verbose > 1:
+                            print(f"File {data.config.save_path} already exists, but will be overwritten (when using run_ACID) due to setting.")
+                    else:
+                        if self.verbose > 0:
+                            print(f"File {data.config.save_path} already exists. The data for this order will be loaded from this file.")
+                        data = Data.load(data.config.save_path) # load the existing data from the file instead of using the newly initialized data
+                else:
+                    data.save() # save the newly initialized, but mostly empty data instance to the file for future reference and use
+
             datalist.append(data) # finally append to the datalist
 
         self.data_list = datalist # datalist property handles the rest
@@ -1717,7 +1777,7 @@ class DataList:
         use_index_mapping : bool                 = True,
         worker            : IntLike|None         = None,
         nworkers          : IntLike|None         = None,
-        allow_overwrite   : bool                 = False,
+        overwrite         : bool|None            = None,
         overwrite_kwargs  : bool                 = False,
         **kwargs,
         ) -> None:
@@ -1748,8 +1808,9 @@ class DataList:
             If the sampler exceeds this size, it will not be stored regardless of the store_sampler flag.
             This is to avoid accidentally storing very large samplers. If None, no limit is set. Default is 1GB.
             A warning will be printed if this size_limit forces the store_sampler to be False if store_sampler was set to True.
-        allow_overwrite : bool, optional
-            If True, will allow overwriting existing result pickles in the save_dir. Default is False, which will skip running ACID on orders 
+        overwrite : bool, optional
+            If True, will allow overwriting existing data and sampler pickles in the save_dir. Default is None, which will use the class
+            default behaviour set in initialization (which is False). If False, this will skip running ACID on orders 
             that already have result pickles in the save_dir.
         overwrite_kwargs : bool, optional
             If True, any keys in the kwargs that are also in the config for the Data instance will be overwritten by the kwargs values.
@@ -1761,6 +1822,10 @@ class DataList:
         """
         from .acid import Acid # local import to avoid circular imports, since Acid imports Data
 
+        # Configure overwrite from class default if not input in the method call
+        if overwrite is None:
+            overwrite = self.overwrite
+
         # Validate worker and nworkers inputs for splitting orders across workers, and set defaults if not provided for easier logic below.
         if worker is not None or nworkers is not None:
             if worker is None or nworkers is None:
@@ -1802,16 +1867,22 @@ class DataList:
         iterable = tqdm(orders, "Running ACID on orders", unit="order") if self.verbose > 1 else orders
         for order in iterable:
 
+            data = self.data_list[self.o2i[order]]
+
             # Check if ACID already ran for this order
-            if os.path.exists(self.data_list[self.o2i[order]].config.save_path) and not allow_overwrite:
-                if self.verbose > 1:
-                    print(f"ACID result for order {order} already exists at {self.data_list[self.o2i[order]].config.save_path}. \n"
-                            f"Skipping this order. To overwrite existing results, set allow_overwrite=True.")
-                # else the sampler and data instance is overwritten
-                continue
+            if os.path.exists(data.config.save_path) and overwrite is False:
+                if data.complete:
+                    if self.verbose > 1:
+                        print(f"An ACID completed result for order {order} already exists. \n"
+                                f"Skipping this order. To overwrite existing results, set overwrite=True.")
+                    continue
+                elif data.exception is not None:
+                    if self.verbose > 1:
+                        print(f"An ACID run for order {order} previously encountered an exception. \n"
+                                f"Skipping this order. To retry and overwrite existing results, set overwrite=True.")
+                    continue
 
-            # Handling if any kwargs were input 
-            data = self.data_list[self.o2i[order]]
+            # Handling if any kwargs were input
             # Only overwrite if overwrite_kwargs is True, otherwise keep the existing linelist/velocities in the Data instance
             if "linelist" in kwargs:
                 ll = kwargs.pop("linelist")
@@ -1866,20 +1937,23 @@ class DataList:
         save_dir : str | None, optional
             The directory to save the DataList pickle file. If None, self.save_dir is used. Default is None.
         """
-        d = {}
         if save_dir is not None:
             self.save_dir = save_dir
         if self.save_dir is None:
-            raise ValueError("No save path provided and save_dir was not set.")
-        d["dict_list"] = [data.to_dict() for data in self.data_list]
+            raise ValueError("No save directory provided and save_dir was not set.")
+        for data in self.data_list:
+            # Ensures that the save paths for each data instance are correct and updated to match the current save_dir,
+            # even if it was changed since initialization.
+            self._set_paths_for_data(data, self.save_dir)
         save_loc = os.path.join(self.save_dir, "datalist.pkl")
-        d["save_dir"] = self.save_dir
+        d = {}
         d["verbose"] = self.verbose
+        # and maybe other class attributes later
         with open(save_loc, "wb") as f:
             pickle.dump(d, f, protocol=pickle.HIGHEST_PROTOCOL)
 
     @classmethod
-    def load(cls, path:str, print_progress:bool=True) -> DataList:
+    def load(cls, path:str) -> DataList:
         """
         Loads a DataList from a pickle file. The pickle file should contain a dictionary with the list of Data objects (converted to dictionaries) and the save_dir.
         Will attempt to load from datalist.pkl in the provided path if it is a directory, otherwise will attempt to load from the provided path directly. 
@@ -1888,55 +1962,55 @@ class DataList:
         Parameters
         ----------
         path : str
-            The path to the pickle file or directory to search for the pickle file.
-        print_progress : bool, optional
-            If True, will print progress messages during the loading process.
-            Similar to verbosity, but only for this function.
-            Default is True.
+            The directory containing the datalist.pkl file, or the datalist.pkl itself. Note that the directories containing the results should also be in here.
 
         Returns
         -------
         DataList
             The loaded DataList object.
         """
-        if os.path.isdir(path):
-            path_check = os.path.join(path, "datalist.pkl")
-            if not os.path.exists(path_check):
-                # Final attempt to directly load the data pickles from order folders
-                all_files = os.listdir(path)
-                data_list = []
-                if print_progress:
-                    all_files = tqdm(all_files, "Opening data pickles in order folders", unit="folder")
-                for folder in all_files:
-                    folder_path = os.path.join(path, folder)
-                    if os.path.isdir(folder_path) and folder.startswith("order_"):
-                        pickle_path = os.path.join(folder_path, "data.pkl")
-                        if os.path.exists(pickle_path):
-                            with open(pickle_path, "rb") as f:
-                                d = pickle.load(f)
-                            data_list.append(Data().from_dict(d))
-                if len(data_list) > 0:
-                    if print_progress:
-                        print(f"Successfully loaded {len(data_list)} Data instances from order folders in {path}.")
-                    return cls.from_datalist(data_list, save_dir=path)
-                else:
-                    raise ValueError(f"No datalist.pkl found in {path}, and no data pickles found in order folders within that path.")
+        abspath = os.path.abspath
+        join = os.path.join
+        isdir = os.path.isdir
+        exists = os.path.exists
+
+        path = abspath(path)
+        if path.endswith("datalist.pkl"):
+            if not exists(path):
+                raise ValueError(f"No pickle file found at {path} to load the DataList from.")
             else:
-                path = path_check
+                path = os.path.dirname(path)
+        elif not isdir(path):
+            raise ValueError(f"The provided path {path} is not a directory, or a datalist pickle file.\n"
+                             f"You should provide a path to a directory containing the folders with the data pickles and sampler files.")
+
+        if exists(join(path, "datalist.pkl")):
+            with open(join(path, "datalist.pkl"), "rb") as f:
+                d = pickle.load(f)
+            verbose = d["verbose"]
         else:
-            if not os.path.exists(path):
-                raise ValueError(f"No pickle file found at {path} to load the DataList from.")
+            verbose = None
+        verbose = Config(verbose=verbose).verbose
 
-        if print_progress:
-            print(f"Loading DataList from {path}...")
-        with open(path, "rb") as f:
-            d = pickle.load(f)
-        data_list = [Data().from_dict(d) for d in d["dict_list"]]
-        verbose = d["verbose"] if "verbose" in d else None
-        # We use a new save_dir depending on the path location in case the directory has changed since last saved
-        obj = cls.from_datalist(data_list, save_dir=os.path.dirname(path), verbose=verbose)
-        if print_progress:
-            print(f"Successfully loaded DataList from {path}.")
+        dir_list = os.listdir(path)        
+        data_list = []
+        folder_moved_flag = False
+        dir_list = dir_list if verbose < 2 else tqdm(dir_list, "Loading Data instances from directory", unit="folder")
+        for folder in dir_list:
+            if isdir(join(path, folder)) and folder.startswith("order_"):
+                save_path = join(path, folder, "data.pkl")
+                sampler_path = join(path, folder, "sampler.h5")
+                if exists(save_path):
+                    data = Data.load(save_path)
+                    if abspath(data.config.save_path) != save_path or abspath(data.config.sampler_path) != sampler_path:
+                        folder_moved_flag = True
+                        cls._set_paths_for_data(data, path)
+                    data_list.append(data)
+
+        if folder_moved_flag and verbose is not None and verbose > 0:
+            print(f"Warning: At least one of the Data instances found in the directory does not match the current location, it has been updated.")
+                    
+        obj = cls.from_datalist(data_list, save_dir=path, verbose=verbose)
         return obj
 
     @property
@@ -1947,10 +2021,10 @@ class DataList:
     def save_dir(self, dir):
         if dir is not None:
             os.makedirs(dir, exist_ok=True)
-        self._save_dir = dir
-        if self._save_dir is None:
-            if self.verbose > 1:
+        elif self._save_dir is None:
+            if self.verbose > 0:
                 print("Warning: save_dir is set to None. No results will be saved. This is not recommended.")
+        self._save_dir = dir
         return
 
     @property
@@ -2075,3 +2149,14 @@ class DataList:
         from .profiles import Profiles
         profiles = Profiles(self.velocities, *self.combined_profile)
         return profiles.plot_fit(**kwargs)
+
+    @staticmethod
+    def _set_paths_for_data(data: Data, save_dir: str) -> None:
+        "Helper to set paths for a Data instance to a new one for a given order."
+        order = data.config.order
+        save_path = os.path.abspath(os.path.join(save_dir, f"order_{order}", "data.pkl"))
+        sampler_path = os.path.abspath(os.path.join(save_dir, f"order_{order}", "sampler.h5"))
+
+        data.config.save_path = save_path
+        data.config.sampler_path = sampler_path
+        data.save()

ACID_code/lsd.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 import numpy as np
 from astropy.io import  fits
-import glob, psutil, os
+import glob, psutil, os, traceback
 import matplotlib.pyplot as plt
 from scipy.signal import find_peaks
-from scipy.interpolate import LSQUnivariateSpline
 from tqdm import tqdm
 from scipy.linalg import cho_factor, cho_solve
 from beartype import beartype
@@ -25,7 +24,7 @@ class LSD:
     def __init__(
             self,
             data    : object|None           = None,
-            OD      : bool                  = True,
+            od      : bool                  = None,
             verbose : IntLike|bool|str|None = None,
             ) -> None:
         """Initialises the LSD class, optionally with a Data instance to take parameters from.
@@ -34,8 +33,9 @@ class LSD:
         ----------
         data : object | None, optional
             A data instance to draw parameters and configs from, by default None
-        OD : bool, optional
-            Whether to perform LSD in optical depth space (True) or flux space (False), by default True.
+        od : bool, optional
+            Whether to perform LSD in optical depth space (True) or flux space (False), by default None.
+            If None, takes from Data instance if provided, else defaults to True.
             We generally recommend always using optical depth as ACID was always intended, but you can set
             this to False if you wish to do your own testing. See :ref:`LSD` in the documentation for more details.
         verbose : :py:type:`IntLike | bool | str | None`, optional
@@ -46,10 +46,10 @@ class LSD:
         # Set class variables, taking from input data if it exists, else setting to defaults
         self.slurm    = "SLURM_JOB_ID" in os.environ
         self.data     = data if data is not None else Data()
-        self.linelist = data.linelist if data is not None else None
-        self.OD       = OD
+        self.linelist = self.data.linelist if self.data is not None else None
+        self.od       = od if od is not None else self.data.config.od
         try:
-            self.config = data.config
+            self.config = self.data.config
         except:
             self.config = Config() # uses defaults
         self.config.update_hipri(verbose=verbose) # Update config with new values, if not None
@@ -87,9 +87,9 @@ class LSD:
             decomposition and solving for the profile, by default None
         """
         # Ensure inputs are numpy arrays
-        wavelengths = np.asarray(wavelengths)
-        flux = np.asarray(flux)
-        errors = np.asarray(errors)
+        wavelengths = np.array(wavelengths)
+        flux = np.array(flux)
+        errors = np.array(errors)
 
         # Ensure dimensions match
         if not wavelengths.shape == flux.shape == errors.shape:
@@ -107,16 +107,23 @@ class LSD:
         # Clip linelist to wavelength range of spectrum
         wavelengths_linelist, depths_linelist = utils.clip_wavelengths(wavelengths, wavelengths_linelist, depths_linelist)
         if len(wavelengths_linelist) == 0:
-            raise LineListRangeError(f"No lines in linelist are within the wavelength range of the observed spectrum. \n"\
-                                     f"You may have mismatched wavelengths units between linelist and spectrum or an empty linelist.\n"\
-                                     f"Please check your linelist and input spectrum.")
+            error = LineListRangeError(
+                "No lines in linelist are within the wavelength range of the observed spectrum.\n"
+                "You may have mismatched wavelength units between linelist and spectrum or an empty linelist.\n"
+                "Please check your linelist and input spectrum."
+            )
+            self.data.exception = error
+            self.data.traceback = traceback.format_stack()
+            raise error
 
         # Apply S/N cut (of 1/(3*SN)) to linelist
         wavelengths_linelist, depths_linelist = self.sn_clip(wavelengths_linelist, depths_linelist, sn)
 
         # Convert to optical depth space for the linelist and the spectrum if needed, and convert errors accordingly
-        if self.OD:
+        if self.od:
             flux, errors, depths_linelist = utils.flux_to_od(flux, errors, depths_linelist)
+        else:
+            flux -= 1
 
         # Calculates alpha in optical depth, selects lines greater than 1/(3*sn)
         if alpha is None:
@@ -130,11 +137,17 @@ class LSD:
         # Solve for profile and profile errors using Cholesky factors
         self.profile, self.profile_errors, self.cov_z = self.solve_z(self.alpha, flux, errors, self.c_factor, return_cov=True)
 
+        self.forward_model = self.alpha @ self.profile
+        self.forward_model_errors = np.sqrt(np.sum((self.alpha * self.profile_errors)**2, axis=1))
+
         # Convert profile back to flux if needed
-        if self.OD:
+        if self.od:
             self.profile_F, self.profile_errors_F, self.cov_z_F = utils.od_to_flux(self.profile, self.profile_errors, cov_matrix=self.cov_z)
+            self.forward_model, self.forward_model_errors = utils.od_to_flux(self.forward_model, self.forward_model_errors)
         else:
+            self.profile += 1
             self.profile_F, self.profile_errors_F, self.cov_z_F = self.profile, self.profile_errors, self.cov_z
+            self.forward_model += 1
 
         return
 
@@ -171,7 +184,10 @@ class LSD:
         nrest = np.sum(idx)
         perc = 100 * nrest / (nrest + ncut)
         if nrest == 0:
-            raise SNCutError(f"No lines remain in the linelist after S/N cut. Please check your linelist and S/N value.")
+            error = SNCutError(f"No lines remain in the linelist after S/N cut. Please check your linelist and S/N value.")
+            self.data.exception = error
+            self.data.traceback = traceback.format_stack()
+            raise error
         if self.config.verbose > 0:
             if perc < 5:
                 print("Warning: Less than 5% of lines remain after S/N cut. Check your linelist and S/N value.")

ACID_code/mcmc.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import numpy as np
+from numpy.linalg import norm
 from . import utils
 from .utils import Array1D, Array2D
 from beartype import beartype
@@ -20,6 +21,12 @@ def _mp_log_probability(theta):
     """Wrapper for log probability function for multiprocessing."""
     return _MCMC(theta)
 
+def _mp_log_likelihood(theta):
+    return _MCMC.dynesty_logprob(theta)
+
+def _mp_ptform(u):
+    return _MCMC.ptform(u)
+
 class MCMC:
 
     """
@@ -38,6 +45,9 @@ class MCMC:
             velocities            : Array1D|None = None,
             c_factor                             = None,
             deterministic_profile : bool         = False,
+            sampler_type          : str          = "emcee",
+            model_inputs          : Array1D|None = None,
+            od                    : bool         = True
         ) -> None:
         """
         Initialise MCMC functions with necessary data.
@@ -75,6 +85,9 @@ class MCMC:
             self.velocities = data.velocities
             self.c_factor = data.c_factor
             self.deterministic_profile = data.config.deterministic_profile
+            self.sampler_type = data.config.sampler_type
+            self.model_inputs = data.model_inputs
+            self.od = data.config.od
         else:
             self.x = x_or_data
             self.y = y
@@ -83,6 +96,10 @@ class MCMC:
             self.velocities = velocities
             self.c_factor = c_factor
             self.deterministic_profile = deterministic_profile
+            self.sampler_type = sampler_type
+            self.model_inputs = model_inputs
+            self.od = od
+            data = None
 
         self.k_max = self.alpha.shape[1] # the number of velocity points in the profile
 
@@ -91,10 +108,15 @@ class MCMC:
         a, b = utils.get_normalisation_coeffs(self.x)
         self.u = (a * self.x) + b # These are the normalized wavelengths used throughout the fitting process
 
-        # For deterministic model, the below variables are used, and are precomputed for speed
-        err_od = self.yerr / self.y # independent of continuum, since it's a ratio
-        V = 1.0 / (err_od ** 2) # variance vector in log space, error already in log space
-        self.AtV = self.alpha.T * V # precompute alpha matrix multiplication for _mcmc_solve_z input
+        if self.od:
+            # For deterministic model, the below variables are used, and are precomputed for speed
+            err_od = self.yerr / self.y # independent of continuum, since it's a ratio
+            V = 1.0 / (err_od ** 2) # variance vector in log space, error already in log space
+        else:
+            # For non-OD case, we need to precompute the variance vector in flux space for the likelihood calculation
+            V = 1.0 / (self.yerr ** 2) # variance vector in flux space
+
+        self.AtV = self.alpha.T * V # precompute alpha matrix multiplication for
 
         # Configure whether to use full or deterministic model
         if self.deterministic_profile is False:
@@ -125,10 +147,11 @@ class MCMC:
         """
         # Extract profile points and continuum coefficients from theta
         z = theta[:self.k_max]
+        z -= 1 if not self.od else 0 # if not using OD, profile points are in flux space and need to be shifted by 1
         mdl = self.alpha @ z
 
         # Converting model from optical depth to flux
-        mdl = np.exp(-mdl)
+        mdl = np.exp(-mdl) if self.od else mdl+1 # if not using OD, just use flux directly
 
         # Calculate continuum polynomial
         coefs = np.asarray(theta[self.k_max:], dtype=float)
@@ -163,13 +186,22 @@ class MCMC:
 
         # Calculate fitted flux and convert to OD
         fitted_flux = self.y/mdl
-        flux_od = - np.log(fitted_flux)
+
+        # Do OD/non-OD conversions
+        if self.od:
+            flux_od = (-np.log(fitted_flux))
+            AtV = self.AtV
+        else:
+            AtV = self.AtV * (mdl * mdl)
+            flux_od = fitted_flux - 1
 
         # Solve for the profile points
-        z = cho_solve(self.c_factor, self.AtV @ flux_od)
+        z = cho_solve(self.c_factor, AtV @ flux_od, check_finite=False)
 
         # Convert back from optical depth to flux
-        forward = np.exp(- (self.alpha @ z)) * mdl
+        dot_prod = self.alpha @ z
+        dot_prod = np.exp(-dot_prod) if self.od else dot_prod + 1
+        forward = dot_prod * mdl
 
         return forward, z
 
@@ -190,8 +222,12 @@ class MCMC:
         """
 
         # Hard box prior on each z[i]
-        if np.any((z < -0.4) | (z > 1.6)):
-            return -np.inf
+        if self.od:
+            if np.any((z < -0.4) | (z > 1.6)):
+                return -np.inf
+        else:
+            if np.any((z > 0.5) | (z <= -1)):
+                return -np.inf
 
         # # excluding the continuum points in the profile (in flux)
         # z_cont = []
@@ -289,4 +325,53 @@ class MCMC:
         tol_str = f"{last_tolerance:.4f}{tol_str}{config.tau_tol}"
         neff_str = ">" if last_neff > config.min_tau_factor else "<"
         neff_str = f"{last_neff:.2f}{neff_str}{config.min_tau_factor}"
-        return tol_str, neff_str
+        return tol_str, neff_str
+
+    def dynesty_logprob(self, theta):
+        """Log likelihood function for dynesty nested sampling."""
+
+        forward, z = self.model_function(theta)
+
+        if not np.all(np.isfinite(forward)):
+            return -np.inf
+
+        lp = self.log_prior(z)
+        if not np.isfinite(lp):
+            return -np.inf
+
+        diff = self.y - forward
+        var = self.yerr * self.yerr
+
+        return -0.5 * np.sum(diff * diff / var + np.log(2 * np.pi * var))
+    
+    def ptform(self, u):
+        """
+        Prior transform for dynesty.
+
+        Maps unit-cube samples u in [0, 1] to continuum polynomial
+        coefficients using uniform priors centred on self.model_inputs.
+        """
+
+        u = np.asarray(u, dtype=float)
+        theta0 = np.asarray(self.model_inputs, dtype=float)
+        theta0 = theta0[self.k_max:] # only continuum coefficents, not profile points
+
+        # Width of uniform prior around curve_fit solution.
+        # The floor matters because higher-order polynomial coefficients
+        # may be close to zero.
+        frac_width = 5
+        abs_floor = 0.05
+
+        width = np.maximum(frac_width * np.abs(theta0), abs_floor)
+
+        # Usually the zeroth-order continuum coefficient is close to 1,
+        # so give it a slightly wider absolute floor.
+        if len(width) > 0:
+            width[0] = max(width[0], 0.25)
+
+        lower = theta0 - width
+        upper = theta0 + width
+
+        return lower + u * (upper - lower)
+
+            

ACID_code/result.py

@@ -2,17 +2,21 @@ from __future__ import annotations
 from time import time
 import numpy as np
 import matplotlib.pyplot as plt
-import corner, sys, os, pickle, warnings, contextlib, functools, inspect, psutil
+import corner, sys, os, warnings, contextlib, functools, inspect
 from emcee import EnsembleSampler
-import emcee.backends.backend as emceebackend
 from beartype import beartype
 from scipy.interpolate import interp1d
 from numpy.polynomial import polynomial as P
 from .lsd import LSD
-from . import mcmc
 from . import utils
 from .data import Data
 from .utils import IntLike, Scalar
+try:
+    from dynesty.sampler import Sampler
+    from dynesty import plotting as dyplot
+except ImportError:
+    Sampler = None
+    dyplot = None
 #TODO: utils.set_dict_defaults for plots
 
 warnings.filterwarnings("ignore")
@@ -59,7 +63,7 @@ class Result:
     def __init__(
             self,
             data                    : Data|object,
-            sampler                 : EnsembleSampler|None  = None,
+            sampler                 : EnsembleSampler|Sampler|None = None, # type:ignore
             process_results         : bool                  = True,
             verbose                 : IntLike|bool|str|None = None,
         ) -> None:
@@ -74,7 +78,7 @@ class Result:
             provided, a sampler can be provided in the second argument. If a sampler object 
             is provided, it will be used as the sampler, but all other attributes will need 
             to be set manually for the Result object to be fully functional.
-        sampler : :py:class:`emcee.EnsembleSampler`, optional
+        sampler : :py:class:`emcee.EnsembleSampler` | :py:class:`dynesty.Sampler`, optional
             Sets and overwrites the sampler in the Data object with this if provided, by default None. 
         process_results : bool, optional
             Whether to process the results from the Acid object upon initialisation, by default True.
@@ -108,7 +112,8 @@ class Result:
         # Handle the sampler if input, initiate if one exists
         self.sampler = sampler if sampler is not None else self.sampler # update sampler if provided, otherwise keep the same
         if self.sampler is not None:
-            self.initiate_sampler(self.sampler) # set internal variables based on sampler, sets sampler_initialized to True
+            self.dynesty = isinstance(self.sampler, Sampler)
+            self.initiate_sampler(self.sampler) # set internal variables based on sampler, sets sampler_initialiated to True
 
         if not self.data.complete:
             if process_results:
@@ -135,7 +140,10 @@ class Result:
         t0 = time()
 
         # Obtain flattened samples
-        flat_samples = self.sampler.get_chain(discard=self.burnin, thin=self.thin, flat=True)
+        if self.dynesty:
+            flat_samples = self.sampler.results.samples_equal()
+        else:
+            flat_samples = self.sampler.get_chain(discard=self.burnin, thin=self.thin, flat=True)
 
         # Getting the final profile and continuum values
         nvel = len(self.data.velocities) if self.config.deterministic_profile is False else 0
@@ -213,15 +221,21 @@ class Result:
             alpha = self.data.alpha if condition else None
 
             LSD_profiles = LSD(self.data)
-            LSD_profiles.run_LSD(wavelengths, flux, error, sn=sn, alpha=alpha)
+            LSD_profiles.run_LSD(wavelengths, flux, error, sn, alpha=alpha)
 
             profile_f = LSD_profiles.profile_F
             profile_errors_f = LSD_profiles.profile_errors_F
             cov_z_f = LSD_profiles.cov_z_F
 
             if counter == 0:
+                # Set combined profile params
                 self.data.combined_profile = [profile_f, profile_errors_f, cov_z_f]
                 self.data.continuum_model = mdl
+
+                # Set the forward model params, multiplied by mdl as LSD is run on normalized flux
+                self.data.forward_model = LSD_profiles.forward_model * mdl
+                self.data.forward_errors = LSD_profiles.forward_model_errors * mdl
+                self.data.forward_x = wavelengths
             else:
                 profiles.append([profile_f, profile_errors_f, cov_z_f])
 
@@ -379,13 +393,23 @@ class Result:
             return fig, ax
         plt.show()
 
+    @_require_sampler
+    def plot_traceplot(self, return_fig:bool=False) -> None | tuple:
+        if not self.dynesty:
+            raise ValueError("Traceplot is only available for dynesty samplers, as emcee traceplots are already plotted in plot_walkers.")
+        fig, ax = dyplot.traceplot(self.sampler.results, labels=self.default_param_labels)
+        plt.suptitle('Dynesty Traceplot')
+        if return_fig:
+            return fig, ax
+        plt.show()
+
     @_require_sampler
     def plot_corner(
         self,
         sampler    :EnsembleSampler|None = None,
         return_fig :bool                 = False,
         **kwargs,
-        ) -> None | tuple:
+        ) -> None | plt.Figure:
         """Creates a corner plot for at maximum the last 8 LSD profile and continuum polynomial coefficients.
 
         Parameters
@@ -402,6 +426,13 @@ class Result:
         ----------
         If return_fig is True, returns the figure object containing the corner plot, else None
         """
+        if self.dynesty:
+            fig, axes = dyplot.cornerplot(self.sampler.results, labels=self.default_param_labels, show_titles=True, title_fmt=".3f", title_kwargs={"fontsize": 16}, **kwargs)
+            plt.suptitle('Dynesty Corner Plot')
+            if return_fig:
+                return fig, axes
+            plt.show()
+            return
 
         # Get samples and thin and burnin from the class variables
         samples = self.sampler.get_chain()
@@ -496,6 +527,7 @@ class Result:
     @_require_profiles
     def plot_forward_model(
         self,
+        fig_ax          :tuple|None = None,
         grid            :bool      = True,
         labels          :dict|None = None,
         return_fig      :bool      = False,
@@ -506,6 +538,12 @@ class Result:
 
         Parameters
         ----------
+        fig_ax: tuple | None
+            Optionally provide an existing fig/axis tuple to plot on, by default None and
+            creates a new figure and axis. The axis must be a 2 element array of axes, 
+            where the first axis is for the spectrum and forward model, 
+            and the second axis is for the residuals.
+            If provided, the grid, labels, and titles should be set by you.
         grid : bool, optional
             Show or hide grid, by default True
         labels : dict | None, optional
@@ -544,8 +582,7 @@ class Result:
 
         # Get flat_samples which are the same samples used to calculate the final profile, alpha is OD, 
         # so convert profile back to OD and reconvert to flux for forward model
-        profile = utils.flux_to_od(self[0])
-        model_flux = utils.od_to_flux(self.data.alpha @ profile) * self.data.continuum_model
+        model_flux = self.data.forward_model
 
         # Due to distortion at the edges of the profile, we drop the last 2 pixels
         wavelengths = utils.drop_edges(wavelengths)
@@ -554,22 +591,26 @@ class Result:
         continuum_model = utils.drop_edges(self.data.continuum_model)
 
         # Plotting
-        fig, ax = plt.subplots(2, 1, **subplot_kwargs)
+        if fig_ax is not None:
+            fig, ax = fig_ax
+        else:
+            fig, ax = plt.subplots(2, 1, **subplot_kwargs)
+            ax[0].set_title(labels["title"])
+            ax[1].set_xlabel(labels["xlabel"])
+            ax[0].set_ylabel(labels["ylabel"])
+            ax[1].set_ylabel(labels["residuals_ylabel"])
+            ax[0].grid(grid)
+            ax[1].grid(grid)
+            plt.subplots_adjust(hspace=0.05)
+
+        ax[1].axhline(0, color='black', linestyle='--', linewidth=1)
         ax[0].plot(wavelengths, flux, color='black', linewidth=1, label='Observed Spectrum')
         ax[0].plot(wavelengths, model_flux, color='C0', linewidth=1, label='Forward Model Fit')
         ax[0].plot(wavelengths, continuum_model, color='C1', linewidth=1, label='Fitted Continuum', linestyle='--')
         ax[1].plot(wavelengths, model_flux-flux, color='C0', linewidth=1, label='Residuals')
         ax[1].axhline(0, color='black', linestyle='--', linewidth=1)
-        ax[0].set_title(labels["title"])
-        ax[1].set_xlabel(labels["xlabel"])
-        ax[0].set_ylabel(labels["ylabel"])
-        ax[1].set_ylabel(labels["residuals_ylabel"])
-        ax[1].axhline(0, color='black', linestyle='--', linewidth=1)
         ax[0].legend()
         ax[1].legend()
-        ax[0].grid(grid)
-        ax[1].grid(grid)
-        plt.subplots_adjust(hspace=0.05)
 
         if return_fig:
             return fig, ax
@@ -720,14 +761,14 @@ class Result:
             return fig, ax
         plt.show()
 
-    def initiate_sampler(self, sampler:EnsembleSampler|None, _method_name=None) -> None:
+    def initiate_sampler(self, sampler:EnsembleSampler|Sampler|None, _method_name=None) -> None: # type:ignore
         """
         Initiates the sampler attribute from an external sampler.
 
         Parameters
         ----------
-        sampler : :py:class:`emcee.EnsembleSampler`
-            An emcee EnsembleSampler object to set as the sampler attribute.
+        sampler : :py:class:`emcee.EnsembleSampler` or object, optional
+            An emcee EnsembleSampler object or a compatible sampler object to set as the sampler attribute.
         _method_name : str, optional
             Internal parameter used to track which method is calling initiate_sampler, for error messages. 
             Not intended for user input, by default None.
@@ -744,6 +785,14 @@ class Result:
                 error_msg = "Cannot initiate sampler without a sampler stored in the instance or passed as a parameter, please pass in a sampler "
             raise AttributeError(error_msg)
 
+        if self.dynesty:
+            a=ord('a')
+            alph=[chr(i) for i in range(a,a+26)]
+            poly_labels = [alph[i] for i in range(self.config.poly_ord + 1)]
+            self.default_param_labels = poly_labels
+            self.default_params = None
+            return
+
         # Calculate autocorr time, burnin, thin
         # Suppress output from get_autocorr_time call
         with open(os.devnull, "w") as devnull, \
@@ -799,12 +848,12 @@ class Result:
         self.sampler_initialized = True
 
     @property
-    def sampler(self) -> EnsembleSampler|None:
+    def sampler(self) -> EnsembleSampler|Sampler|None: # type:ignore
         """Returns the sampler attribute, by default is None if not saved."""
         return self.data.sampler
 
     @sampler.setter
-    def sampler(self, value: EnsembleSampler|None) -> None:
+    def sampler(self, value: EnsembleSampler|Sampler|None) -> None: # type:ignore
         """Sets the sampler in the data class."""
         self.data.sampler = value
 

ACID_code/utils.py

@@ -108,9 +108,12 @@ def mask_invalid(wavelengths, flux, errors=None, return_mask=False, verbose=2):
 
     if verbose > 1:
         num_invalid = np.size(wavelengths) - np.count_nonzero(mask)
-        if num_invalid > 0:
-            print(f"Your spectrum includes {num_invalid} out of {np.size(wavelengths)} non-positive/non-finite/nan values, which will be dropped when necessary, \n"
-                  f"but it is still recommended to check your wavelength, spectrum and error arrays for bad pixels and make sure this is intentional.")
+        perc_invalid = num_invalid / np.size(wavelengths) * 100
+        if perc_invalid > 10:
+            print(f"Your spectrum includes {num_invalid} out of {np.size(wavelengths)} non-positive/non-finite/nan values ({perc_invalid:.2f}%), \n"
+                  f"which will be dropped when necessary, but it is still recommended to check your wavelength, \n"
+                  f"spectrum and error arrays for bad pixels and make sure this is intentional. \n"
+                  f"This warning is only printed if more than 10% of pixels are invalid.")
 
     output = (w, f, e) if errors is not None else (w, f)
     output = output + (mask,) if return_mask else output
@@ -146,8 +149,9 @@ def drop_invalid(wavelengths, flux, errors=None, return_mask=False, verbose=2):
 
     if verbose > 1:
         num_invalid = np.size(wavelengths) - np.count_nonzero(mask)
-        if num_invalid > 0:
-            print(f"Dropped {num_invalid} invalid pixels out of {np.size(wavelengths)} (non-finite or <= 0).")
+        perc_invalid = num_invalid / np.size(wavelengths) * 100
+        if perc_invalid > 10:
+            print(f"Dropped {num_invalid} invalid pixels out of {np.size(wavelengths)} (non-finite or <= 0), which is {perc_invalid:.2f}% of the total.")
 
     output = (w, f, e) if errors is not None else (w, f)
     output = output + (mask,) if return_mask else output

{acid_code-2.0.0a2.dist-info → acid_code-2.0.0a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ACID_code
-Version: 2.0.0a2
+Version: 2.0.0a3
 Summary: Returns line profiles from input spectra by fitting the stellar continuum and performing LSD
 Author: Lucy Dolan
 Author-email: Benjamin Cadell <bcadell01@qub.ac.uk>

acid_code-2.0.0a3.dist-info/RECORD

@@ -0,0 +1,15 @@
+ACID_code/__init__.py,sha256=-cmUx-kDheDBhk8ZlleWVRrWZyXZRVe-dwciVV9ZxLg,655
+ACID_code/acid.py,sha256=hibDm91O0_HLO3d50I2Q137rXi_xmXdcJSb9KcDNLCs,75532
+ACID_code/data.py,sha256=qaLGYFDPyK7XlXg09b1AZO4xGMZCmrYd7t04SU3LknA,114575
+ACID_code/errors.py,sha256=qqG44x_rVpi5njLIoRkwvw_6owy33wejvsFuAZ-XKKM,525
+ACID_code/load.py,sha256=3gzIZpAv7flgX5ekWdRDI95hkPTxvtp2F1liygpivOQ,5453
+ACID_code/lsd.py,sha256=VfdB-mN2uH20N6jT6lGKhEaepOwB2bJQepyZfIpqM9U,21521
+ACID_code/mcmc.py,sha256=PjGuxUhq-oFRkKKQmQPdSp82NyueG9FPwr7kpLgLk0U,14314
+ACID_code/profiles.py,sha256=KkVsJjCLXx_wjRBialfmnlGwjWsROy0GkggdPPYLelc,17103
+ACID_code/result.py,sha256=itDNK0zxgtNH8-Igqip-zFi6StK5ZyuMS47UYVooSq8,42141
+ACID_code/utils.py,sha256=Xt24WrhV6bKW-NiNl1FsNiJ2krntCJeCUL-hzIQAdLA,27637
+acid_code-2.0.0a3.dist-info/licenses/LICENSE,sha256=L6dUgqjvHmRoobrBCPSHKC4UtRM5Ldp1DJBC4bnLk3w,1070
+acid_code-2.0.0a3.dist-info/METADATA,sha256=XFBGmLrD4od1XTRY4SRGN3Cg-RAvM21VUEClyV6paIQ,2999
+acid_code-2.0.0a3.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+acid_code-2.0.0a3.dist-info/top_level.txt,sha256=O4OaSabv1ebFYQmHgftr1PGAv6BvC2l81Y3HjgNehQI,10
+acid_code-2.0.0a3.dist-info/RECORD,,

acid_code-2.0.0a2.dist-info/RECORD

@@ -1,15 +0,0 @@
-ACID_code/__init__.py,sha256=-cmUx-kDheDBhk8ZlleWVRrWZyXZRVe-dwciVV9ZxLg,655
-ACID_code/acid.py,sha256=Id3pDMmuQ5OEx3gLarS5p5-pnAiZ0b_tF7_DX3Ou8Wo,72704
-ACID_code/data.py,sha256=J8Nz7BAzltxclElWH1iI3rx0yUPDRSjcSCH3q1xjsRE,109492
-ACID_code/errors.py,sha256=qqG44x_rVpi5njLIoRkwvw_6owy33wejvsFuAZ-XKKM,525
-ACID_code/load.py,sha256=3gzIZpAv7flgX5ekWdRDI95hkPTxvtp2F1liygpivOQ,5453
-ACID_code/lsd.py,sha256=iSNjxga-ic1tDCEpEOr5ISSxOsg6GLL_BYqTJi58MTA,20814
-ACID_code/mcmc.py,sha256=kNW1Aj5jfWb5xa95yKJU4NBOp1hnGV3KzHV5DrCtKU8,11434
-ACID_code/profiles.py,sha256=KkVsJjCLXx_wjRBialfmnlGwjWsROy0GkggdPPYLelc,17103
-ACID_code/result.py,sha256=9VXQkovQflDSUIYjwHCAo_cFMs7981IOuKDwlFv8_9w,39819
-ACID_code/utils.py,sha256=SY7qlIM15h8yYz7G_-DqXpYhralFbE0ycpOPvnw6DZw,27323
-acid_code-2.0.0a2.dist-info/licenses/LICENSE,sha256=L6dUgqjvHmRoobrBCPSHKC4UtRM5Ldp1DJBC4bnLk3w,1070
-acid_code-2.0.0a2.dist-info/METADATA,sha256=qVMgFV_XyFpmUQgGZ6zG3gSZXZR4aT7Lbhfdp6mIabA,2999
-acid_code-2.0.0a2.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-acid_code-2.0.0a2.dist-info/top_level.txt,sha256=O4OaSabv1ebFYQmHgftr1PGAv6BvC2l81Y3HjgNehQI,10
-acid_code-2.0.0a2.dist-info/RECORD,,