WeightedGroupBy

.github/workflows/CI.yaml

@@ -44,7 +44,6 @@ jobs:
       - name: Upgrade pip and install doc requirements
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pip-tools
           python -m pip install -e ".[extras,docs]"
       - name: build documentation
         run: |

README.rst

@@ -2,7 +2,7 @@
 anesthetic: nested sampling post-processing
 ===========================================
 :Authors: Will Handley and Lukas Hergt
-:Version: 2.0.0-beta.25
+:Version: 2.0.0-beta.26
 :Homepage: https://github.com/williamjameshandley/anesthetic
 :Documentation: http://anesthetic.readthedocs.io/
 
@@ -191,8 +191,8 @@ Why create another one? In general, any dedicated user of software will find tha
 
 .. code:: python
 
-    from anesthetic import MCMCSamples
-    samples = MCMCSamples(root=file_root)                         # Load the samples
+    from anesthetic import read_chains
+    samples = read_chains(file_root)                              # Load the samples
     samples['omegab'] = samples.omegabh2/(samples.H0/100)**2      # Define omegab
     samples.tex['omegab'] = '$\Omega_b$'                          # Label omegab
     samples.plot_1d('omegab')                                     # Simple 1D plot

anesthetic/_version.py

@@ -1 +1 @@
-__version__ = '2.0.0b25'
+__version__ = '2.0.0b26'

anesthetic/samples.py

@@ -515,7 +515,7 @@ def remove_burn_in(self, burn_in, reset_index=False, inplace=False):
             Indicates whether to modify the existing array or return a copy.
 
         """
-        chains = self.groupby(('chain', '$n_\\mathrm{chain}$'),
+        chains = self.groupby(('chain', '$n_\\mathrm{chain}$'), sort=False,
                               group_keys=False)
         nchains = chains.ngroups
         if isinstance(burn_in, (int, float)):
@@ -574,25 +574,32 @@ def Gelman_Rubin(self, params=None):
                       and 'logL' not in key
                       and 'chain' not in key]
         chains = self[params+['chain']].groupby(
-                ('chain', '$n_\\mathrm{chain}$')
+                ('chain', '$n_\\mathrm{chain}$'), sort=False,
         )
+        nchains = chains.ngroups
 
         # Within chain variance ``W``
         # (average variance within each chain):
-        W = chains.cov().groupby(level=['params', 'labels']).mean().to_numpy()
-        # TODO: the above line should be a weighted mean
-        # --> need to fix groupby for WeightedDataFrames!
-
+        W = chains.cov().groupby(level=('params', 'labels'), sort=False).mean()
         # Between-chain variance ``B``
-        # (variance of the chain means compared to the full mean):
-        means_diff = (chains.mean() - self[params].mean()).to_numpy()
-        B = (means_diff.T @ means_diff) / (chains.ngroups - 1)
-        # B = chains.mean().cov().to_numpy()
-        # TODO: fix once groupby is fixed
-
-        L = np.linalg.cholesky(W)
-        invL = np.linalg.inv(L)
-        D = np.linalg.eigvalsh(invL @ B @ invL.T)
+        # (variance of the chain means):
+        B = np.atleast_2d(np.cov(chains.mean().T, ddof=1))
+        # We don't weight `B` with the effective number of samples (sum of the
+        # weights), here, because we want to notice outliers from shorter
+        # chains.
+        # In order to be conservative, we generally want to underestimate `W`
+        # and overestimate `B`, since `W` goes in the denominator and `B` in
+        # the numerator of the Gelman--Rubin statistic `Rminus1`.
+
+        try:
+            invL = np.linalg.inv(np.linalg.cholesky(W))
+        except np.linalg.LinAlgError as e:
+            raise np.linalg.LinAlgError(
+                "Make sure you do not have linearly dependent parameters, "
+                "e.g. having both `As` and `A=1e9*As` causes trouble.") from e
+        D = np.linalg.eigvalsh(invL @ ((nchains+1)/nchains * B) @ invL.T)
+        # The factor of `(nchains+1)/nchains` accounts for the additional
+        # uncertainty from using a finite number of chains.
         Rminus1 = np.max(np.abs(D))
         return Rminus1
 

anesthetic/utils.py

@@ -530,8 +530,9 @@ class to adjust
     """
     for key, val in cls.__dict__.items():
         doc = inspect.getdoc(val)
-        newdoc = re.sub(pattern, repl, doc, *args, **kwargs)
-        try:
-            cls.__dict__[key].__doc__ = newdoc
-        except AttributeError:
-            pass
+        if doc is not None:
+            newdoc = re.sub(pattern, repl, doc, *args, **kwargs)
+            try:
+                cls.__dict__[key].__doc__ = newdoc
+            except AttributeError:
+                pass