fix warning

anesthetic/read/dnest4.py

@@ -24,7 +24,9 @@ def _determine_columns_and_labels(n_params, header, delim=' '):
 def read_dnest4(root,
                 levels_file='levels.txt',
                 sample_file='sample.txt',
-                sample_info_file='sample_info.txt'):
+                sample_info_file='sample_info.txt',
+                *args,
+                **kwargs):
     """
     Read dnest4 output files.
 
@@ -57,16 +59,22 @@ def read_dnest4(root,
         header = f.readline()
 
     n_params = samples.shape[1]
-    columns, labels = _determine_columns_and_labels(n_params, header)
 
     sample_level = sample_info[:, 0].astype(int)
     logL = sample_info[:, 1]
     logL_birth = levels[sample_level, 1]
 
+    kwargs['label'] = kwargs.get('label', os.path.basename(root))
+    columns, labels = _determine_columns_and_labels(n_params, header)
+    columns = kwargs.pop('columns', columns)
+    labels = kwargs.pop('labels', labels)
+
     return DiffusiveNestedSamples(sample_info=sample_info,
                                   levels=levels,
                                   samples=samples,
                                   logL=logL,
                                   logL_birth=logL_birth,
                                   columns=columns,
-                                  labels=labels)
+                                  labels=labels,
+                                  *args,
+                                  **kwargs)

anesthetic/samples.py

@@ -1434,19 +1434,21 @@ def __init__(self, sample_info, levels, samples, *args, **kwargs):
         self.logzero = kwargs.get('logzero', -1e300)
         sample_info_columns = [kwargs.get('columns') +
                                ['level', 'logL', 'tiebreaker', 'ID']]
-        self.samples_with_info = pandas.DataFrame(
+        self.attrs["samples_with_info"] = pandas.DataFrame(
             np.concatenate([samples, sample_info], axis=1),
             columns=sample_info_columns)
-        self.samples_with_info[['level', 'ID']] = \
-            self.samples_with_info[['level', 'ID']].astype(int)
         level_columns = ['log_X',
                          'log_likelihood',
                          'tiebreaker',
                          'accepts',
                          'tries',
                          'exceeds',
                          'visits']
-        self.levels = pandas.DataFrame(levels, columns=level_columns)
+        self.attrs["levels"] = pandas.DataFrame(levels, columns=level_columns)
+
+    @property
+    def _constructor(self):
+        return NestedSamples
 
     def samples_at_level(self, level_index, label):
         """
@@ -1463,8 +1465,9 @@ def samples_at_level(self, level_index, label):
         -------
             numpy.ndarray:
         """
-        selection = (self.samples_with_info.level == level_index).squeeze()
-        return self.samples_with_info[selection][label].to_numpy()
+        samples = self.attrs["samples_with_info"]
+        selection = (samples.level == level_index).squeeze()
+        return samples[selection][label].to_numpy()
 
     def n_live(self, *args):
         """
@@ -1485,24 +1488,6 @@ def n_live(self, *args):
         """
         return self.num_particles
 
-    def LX(self, beta, logX):
-        """
-        Get LX, e.g., for Higson plot.
-
-        Parameters
-        ----------
-        beta: float
-            temperature
-        logX: np.ndarray
-            prior volumes
-
-        Returns
-        -------
-         LX: np.ndarray
-        """
-        LX = super().LX(beta, logX)
-        return LX
-
     def plot_types(self):
         """
         Get types of plots supported by this class.
@@ -1536,9 +1521,9 @@ def points_to_plot(self, plot_type, label, evolution, beta, base_color):
         List[tuple[float]: colors to use
         """
         if plot_type == 'visited points':
-            base_color = 'C0'
             logX = self.logX().to_numpy()[evolution]
-            levels_to_plot = self.levels[self.levels['log_X'] >= logX]
+            levels = self.attrs["levels"]
+            levels_to_plot = levels[levels['log_X'] >= logX]
             max_level_index = levels_to_plot.tail(1).index.to_list()[0]
             colors = [basic_cmap(base_color)(float(j) / (max_level_index + 1))
                       for j in range(1, max_level_index + 2)]