Changes to plot multimodal posteriors when the option is enabled. To …

xpsi/PostProcessing/_runs.py

@@ -95,13 +95,25 @@ def likelihood(self, obj):
 
     @classmethod
     def load_runs(cls, ID, run_IDs, roots, base_dirs, use_nestcheck,
-                  likelihood=None, **kwargs):
+                  likelihood=None,multi_mode=False,mode_label="mode", **kwargs):
         """ Construct a :class:`~.Runs` instance by loading distinct runs.
 
+        param bool multi_mode:
+            Load the runs so that they are separated in different modes found
+            by a MultiNest run with mmodal=True. Works at the moment only when
+            loading a single run and when having models with the same number of
+            modes.
+
+        param string mode_label:
+            Select a label for the modes to be shown in the plots.
+
         The kwargs will be shared by nested sampling runs. The args must be
         lists that will be zipped to instantiate a set of run backends.
 
         """
+        mode_label ="_"+mode_label
+        cls.mode_label = mode_label
+        cls.multi_mode = multi_mode
         # if there is a transform method available, try to wrap it
         # so that error due to a mismatch in parameter order is bypassed
         if likelihood is not None:
@@ -122,10 +134,64 @@ def _transform(q, **kwargs):
             _transform = None
             _overwrite = False
 
+        if multi_mode:
+            if len(roots)>1:
+                raise Exception("Error: multi_mode=True is not yet working when loading multiple runs.")
+            try:
+                cls.mode_len
+            except:
+                pass
+            else:
+                print("Warning: A run/model using multi_mode=True was previously already loaded. This load will overwrite the number of modes found, and can result in an error if trying to plot the previously loaded results.")
+            for vec in range(len(roots)):
+                filerootpath =_os.path.join(base_dirs[vec], roots[vec])
+                # Read the file and manually handle blank lines
+                with open(filerootpath+"post_separate.dat", 'r') as file:
+                    lines = file.readlines()
+
+                # Convert lines to a numpy array, replacing blank lines with a marker (e.g., None or NaN)
+                data = []
+                for line in lines:
+                    stripped = line.strip()
+                    if stripped:  # Non-blank line
+                        data.append(_np.array(stripped.split(), dtype=float))
+                    else:  # Blank line
+                        data.append(None)
+
+                # Separate modes based on None markers
+                modes = []
+                current_mode = []
+                for row in data:
+                    if row is None:
+                        if current_mode:
+                            modes.append(_np.array(current_mode))
+                            current_mode = []
+                    else:
+                        current_mode.append(row)
+
+                # Append the last mode if not empty
+                if current_mode:
+                    modes.append(_np.array(current_mode))
+                for mode in range(len(modes)):
+                    _np.savetxt(filerootpath+f"mode{mode+1}.txt", modes[mode])
+                    roots.append(roots[vec]+f"mode{mode+1}")
+                    run_IDs.append(run_IDs[vec]+f"{mode_label} {mode+1}")
+                    base_dirs.append(base_dirs[vec])
+                    use_nestcheck.append(use_nestcheck[vec])
+            # Forget about the default xpsi loaded file
+            roots = roots[1:]
+            run_IDs = run_IDs[1:]
+            base_dirs = base_dirs[1:]
+            use_nestcheck = use_nestcheck[1:]
+
+            cls.mode_len =len(modes)
+
+
         runs = []
         for root, run_ID, base_dir, check in zip(roots, run_IDs,
                                                  base_dirs,
                                                  use_nestcheck):
+
             runs.append(NestedBackend(root, base_dir,
                                       ID=run_ID,
                                       use_nestcheck=check,
@@ -143,6 +209,16 @@ def set_subset(self, IDs=None, combine=False, combine_all=False,
         if IDs is None:
             self._subset = self._runs
         else:
+            if self.multi_mode:
+                old_IDs=IDs
+                IDs=[]
+                for ID in  old_IDs:
+                    if self.mode_label in ID:
+                        IDs = old_IDs
+                    else:
+                        root_ID=ID
+                        for mode in range(self.mode_len):
+                            IDs.append(root_ID + f"{self.mode_label} {mode+1}")
             self._subset = [self[ID] for ID in IDs]
 
         if combine and force_combine: # create new run object