Repository: sweights/sweights (language: python)

Issue:
Add python linting
Add some python linting checks and pre-commit checks


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index fdf9e96..c09dd1e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,5 @@ dist/
 src/sweights.egg-info
 **__pycache__
 **.DS_Store
+**.swp
+**.swo
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..eaecef1
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,68 @@
+# To use:
+#
+#     pre-commit run -a
+#
+# Or:
+#
+#     pre-commit install  # (runs every time you commit in git)
+#
+# To update this file:
+#
+#     pre-commit autoupdate
+#
+# See https://github.com/pre-commit/pre-commit
+
+repos:
+# Standard hooks
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.0.1
+  hooks:
+  - id: check-case-conflict
+  - id: check-docstring-first
+  - id: check-executables-have-shebangs
+  - id: check-merge-conflict
+  - id: check-symlinks
+  - id: check-yaml
+  - id: debug-statements
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+  - id: sort-simple-yaml
+  - id: file-contents-sorter
+  - id: trailing-whitespace
+    exclude: ^doc/_static/.*.svg
+
+# Python linter (Flake8)
+- repo: https://gitlab.com/pycqa/flake8
+  rev: 3.9.2
+  hooks:
+  - id: flake8
+    args:
+    - "--max-line-length=88"
+
+# Python formatting
+- repo: https://github.com/psf/black
+  rev: 21.12b0
+  hooks:
+  - id: black
+
+# Python docstring formatting
+- repo: https://github.com/pycqa/pydocstyle
+  rev: 6.1.1
+  hooks:
+  - id: pydocstyle
+    files: src/sweights/[^_].*\.py
+
+# C++ formatting
+- repo: https://github.com/pre-commit/mirrors-clang-format
+  rev: v13.0.0
+  hooks:
+  - id: clang-format
+
+# CMake formatting
+- repo: https://github.com/cheshirekow/cmake-format-precommit
+  rev: v0.6.13
+  hooks:
+  - id: cmake-format
+    additional_dependencies: [pyyaml]
+    types: [file]
+    files: (\.cmake|CMakeLists.txt)(.in)?$
diff --git a/README.md b/README.md
index 4714117..9f55e21 100644
--- a/README.md
+++ b/README.md
@@ -65,4 +65,3 @@ This does the following:
 5. **Fit the weighted distributions and correct the covariance** using the `cov_correct` function provided
 
   ![tfit](https://user-images.githubusercontent.com/1140576/142237505-11032b1c-b6fa-47dc-9a0e-e965210fdf6b.png)
-
diff --git a/doc/about.rst b/doc/about.rst
index 3eae81b..b5ca693 100644
--- a/doc/about.rst
+++ b/doc/about.rst
@@ -16,4 +16,3 @@ Maintainers
 -----------
 
 * Matthew Kenzie (@matthewkenzie)
-
diff --git a/doc/conf.py b/doc/conf.py
index b23a4f5..cb20a57 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -5,6 +5,7 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html
 
 import os
+from sweights import __version__
 
 # -- Path setup --------------------------------------------------------------
 
@@ -19,13 +20,13 @@
 
 # -- Project information -----------------------------------------------------
 
-project = 'sweights'
-copyright = '2021, Matthew Kenzie'
-author = 'Matthew Kenzie'
+project = "sweights"
+copyright = "2021, Matthew Kenzie"
+author = "Matthew Kenzie"
 
 # The full version, including alpha/beta/rc tags
-from sweights import __version__
-release = f'{__version__}'
+
+release = f"{__version__}"
 
 # http://read-the-docs.readthedocs.org/en/latest/theme.html#how-do-i-use-this-locally-and-on-read-the-docs
 # on_rtd is whether we are on readthedocs.org
@@ -54,7 +55,7 @@
 autosummary_generate = True
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # The suffix of source filenames.
 source_suffix = ".rst"
@@ -68,9 +69,9 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ['_build', '_themes', 'Thumbs.db', '.DS_Store']
+exclude_patterns = ["_build", "_themes", "Thumbs.db", ".DS_Store"]
 
-# The reST default role (used for this markup: `text`) to use for all documents.
+# The reST default role (used for this markup: `text`) to use for all documents
 # default_role = None
 
 # If true, '()' will be appended to :func: etc. cross-reference text.
@@ -94,7 +95,7 @@
 # -- Options for HTML output -------------------------------------------------
 
 # autodoc_member_order = "bysource"
-#autodoc_mock_imports = ["scipy"]
+# autodoc_mock_imports = ["scipy"]
 
 # html_logo = "_static/sweights_logo.svg"
 
@@ -106,12 +107,12 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-#html_theme = 'alaaster'
+# html_theme = 'alaaster'
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ["_static"]
 
 # If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
 # using the given strftime format.
@@ -157,7 +158,7 @@
 # Output file base name for HTML help builder.
 htmlhelp_basename = "sweightsdoc"
 
-# -- Options for LaTeX output --------------------------------------------------
+# -- Options for LaTeX output -------------------------------------------------
 
 latex_elements = {
     # The paper size ('letterpaper' or 'a4paper').
@@ -169,9 +170,15 @@
 }
 
 # Grouping the document tree into LaTeX files. List of tuples
-# (source start file, target name, title, author, documentclass [howto/manual]).
+# (source start file, target name, title, author, documentclass [howto/manual])
 latex_documents = [
-    ("index", "latex.tex", "sweights Documentation", "Matthew Kenzie", "manual")
+    (
+        "index",
+        "latex.tex",
+        "sweights Documentation",
+        "Matthew Kenzie",
+        "manual"
+    )
 ]
 
 # The name of an image file (relative to this directory) to place at the top of
@@ -195,17 +202,25 @@
 # latex_domain_indices = True
 
 
-# -- Options for manual page output --------------------------------------------
+# -- Options for manual page output -------------------------------------------
 
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
-man_pages = [("index", "sweights", "sweights Documentation", ["Matthew Kenzie"], 1)]
+man_pages = [
+    (
+        "index",
+        "sweights",
+        "sweights Documentation",
+        ["Matthew Kenzie"],
+        1
+    )
+]
 
 # If true, show URL addresses after external links.
 # man_show_urls = False
 
 
-# -- Options for Texinfo output ------------------------------------------------
+# -- Options for Texinfo output -----------------------------------------------
 
 # Grouping the document tree into Texinfo files. List of tuples
 # (source start file, target name, title, author,
diff --git a/doc/reference.rst b/doc/reference.rst
index 6d3666b..9b2f3e7 100644
--- a/doc/reference.rst
+++ b/doc/reference.rst
@@ -37,4 +37,3 @@ Utilities
 
 .. autofunction:: kendall_tau
 .. autofunction:: convertRooAbsPdf
-
diff --git a/doc/requirements.txt b/doc/requirements.txt
index 9953644..3c89565 100644
--- a/doc/requirements.txt
+++ b/doc/requirements.txt
@@ -4,4 +4,3 @@ matplotlib==3.5.0
 pillow==8.3.2
 sweights==0.0.4
 nbsphinx
-
diff --git a/src/sweights/__init__.py b/src/sweights/__init__.py
index fda1905..7abb07c 100644
--- a/src/sweights/__init__.py
+++ b/src/sweights/__init__.py
@@ -1,4 +1,19 @@
-__version__ = '0.0.4'
+"""Python interface to produce sweights and cows.
+
+* Code: https://github.com/matthewkenzie/sweights
+* Docs: https://sweights.readthedocs.io
+
+"""
+__version__ = "0.0.4"
+__all__ = [
+    "sweight",
+    "convertRooAbsPdf",
+    "cow",
+    "kendall_tau",
+    "cov_correct",
+    "approx_cov_correct",
+]
+
 from sweights.sweight import sweight, convertRooAbsPdf
 from sweights.cow import cow
 from sweights.independence import kendall_tau
diff --git a/src/sweights/covariance.py b/src/sweights/covariance.py
index 21a33a3..5ff8da8 100644
--- a/src/sweights/covariance.py
+++ b/src/sweights/covariance.py
@@ -1,32 +1,32 @@
-# implements a covariance correction for weighted data fits
+"""Implementation of covariance correction for weighted fits."""
+
 import numpy as np
 from scipy.misc import derivative
 
-## derivative of function pdf with respect to variable at index var
-## evaluated at point point
-def partial_derivative(pdf, var, point, data):
-  args = point[:]
-  def wraps(x):
-    args[var] = x
-    return pdf(data,*args)
-  return derivative(wraps,point[var], dx=1e-6)
-
-# you should pass the pdf function
-# which must be in the form pdf(data,*pars) e.g. a 1D Gaussian would be pdf(x,mean,sigma)
-# then pass the data (should be an appropriate degree numpy array to be passed to your pdf
-# then pass the weights (should have same shape as data)
-# then pass the fit values and fitted covariance of the nominal fit
+# derivative of function pdf with respect to variable at index var
+# evaluated at point point
+
+
+def _partial_derivative(pdf, var, point, data):
+    args = point[:]
+
+    def wraps(x):
+        args[var] = x
+        return pdf(data, *args)
+
+    return derivative(wraps, point[var], dx=1e-6)
+
 
 def approx_cov_correct(pdf, data, wts, fvals, fcov, verbose=False):
     """
-    Perform a first order covariance correction for a fit to weighted data
+    Perform a first order covariance correction for a fit to weighted data.
 
     Parameters
     ----------
     pdf : callable
-        The control variable pdf which must take arguments (x, p0,...,pn) which has
-        been fitted to a weighted dataset, where x is the observable and p0 ... pn
-        are the shape parameters
+        The control variable pdf which must take arguments (x, p0,...,pn)
+        which has been fitted to a weighted dataset, where x is the observable
+        and p0 ... pn are the shape parameters
     data : ndarray
         The data values of the observable at which the pdf is evaluated
     wts : ndarray
@@ -34,8 +34,8 @@ def approx_cov_correct(pdf, data, wts, fvals, fcov, verbose=False):
     fvals : array_like
         A list of the fitted values of the shape parameters p0,....,pn
     fcov : ndarray
-        A covariance matrix of the weighted likelihood fit before the correction (this
-        is normally available from the minmiser e.g. iminuit)
+        A covariance matrix of the weighted likelihood fit before the
+        correction (this is normally available from the minmiser e.g. iminuit)
     verbose : bool, optional
         Print some output
 
@@ -47,58 +47,65 @@ def approx_cov_correct(pdf, data, wts, fvals, fcov, verbose=False):
     Notes
     -----
     This function corresponds to the weighted score function with independent
-    weights (first term of Eq.51 in `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_).
-    This is often a good approximation for sweights although will be an overestimate.
-    For a better correction for sweights use `cov_correct` (which is only in general accurate
-    when the shape parameters in the discriminating variable are known)
+    weights (first term of Eq.51 in
+    `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_).
+    This is often a good approximation for sweights although will be an
+    overestimate. For a better correction for sweights use `cov_correct` (which
+    is only in general accurate when the shape parameters in the discriminating
+    variable are known)
 
     See Also
     --------
     cov_correct
     """
-
     dim = len(fvals)
-    assert(fcov.shape[0]==dim and fcov.shape[1]==dim)
+    assert fcov.shape[0] == dim and fcov.shape[1] == dim
 
     Djk = np.zeros(fcov.shape)
 
-    prob = pdf(data,*fvals)
+    prob = pdf(data, *fvals)
 
     for j in range(dim):
-        derivj = partial_derivative(pdf,j,fvals,data)
+        derivj = _partial_derivative(pdf, j, fvals, data)
         for k in range(dim):
-            derivk = partial_derivative(pdf,k,fvals,data)
+            derivk = _partial_derivative(pdf, k, fvals, data)
 
-            Djk[j,k] = np.sum( wts**2 * (derivj*derivk) / prob**2 )
+            Djk[j, k] = np.sum(wts ** 2 * (derivj * derivk) / prob ** 2)
 
     corr_cov = fcov * Djk * fcov.T
 
     if verbose:
-        print('First order covariance correction for weighted events')
-        print('  Original covariance:')
-        print('\t', str(fcov).replace('\n','\n\t '))
-        print('  Corrected covariance:')
-        print('\t', str(corr_cov).replace('\n','\n\t '))
+        print("First order covariance correction for weighted events")
+        print("  Original covariance:")
+        print("\t", str(fcov).replace("\n", "\n\t "))
+        print("  Corrected covariance:")
+        print("\t", str(corr_cov).replace("\n", "\n\t "))
 
     return corr_cov
 
-def cov_correct(hs, gxs, hdata, gdata, weights, Nxs, fvals, fcov, dw_dW_fs, Wvals, verbose=False):
+
+def cov_correct(
+    hs, gxs, hdata, gdata, weights, Nxs, fvals, fcov, dw_dW_fs, Wvals, verbose=False
+):
     """
-    Perform a second order covariance correction for a fit to weighted data
+    Perform a second order covariance correction for a fit to weighted data.
 
     Parameters
     ----------
     hs : callable
-        The control variable pdf which must take arguments (x, p0,...,pn) which has
-        been fitted to a weighted dataset, where x is the observable and p0 ... pn
-        are the shape parameters
+        The control variable pdf which must take arguments (x, p0,...,pn)
+        which has been fitted to a weighted dataset, where x is the observable
+        and p0 ... pn are the shape parameters
     gxs : list of callable
-        A list of the disciminant variable pdfs which must take a single argument (x)
-        where x is the observable (shape parameters of gxs must be known in this case)
+        A list of the disciminant variable pdfs which must take a single
+        argument (x), where x is the observable (shape parameters of gxs must
+        be known in this case)
     hdata : ndarray
-        The data values of the control variable observable at which the `hs` pdf is evaluated
+        The data values of the control variable observable at which the `hs`
+        pdf is evaluated
     gdata : ndarray
-        The data values of the disciminant variable observable at which the `gxs` pdfs are evaluated
+        The data values of the disciminant variable observable at which the
+        `gxs` pdfs are evaluated
     weights : ndarray
         The values of the weights
     Nxs : list of float or tuple of float
@@ -106,12 +113,12 @@ def cov_correct(hs, gxs, hdata, gdata, weights, Nxs, fvals, fcov, dw_dW_fs, Wval
     fvals : array_like
         A list of the fitted values of the shape parameters p0,....,pn
     fcov : ndarray
-        A covariance matrix of the weighted likelihood fit before the correction (this
-        is normally available from the minmiser e.g. iminuit)
+        A covariance matrix of the weighted likelihood fit before the
+        correction (this is normally available from the minmiser e.g. iminuit)
     dw_dW_fs : list of callable
-        A list of the functions describing the partial derivate of the weight function with
-        respect to the W matrix elements (see the `tests/example.py` tutorial to see this
-        passed for sweights and COWs)
+        A list of the functions describing the partial derivate of the weight
+        function with respect to the W matrix elements (see the tutorial to see
+        this passed for sweights and COWs)
     Wvals : list of float or tuple of float
         A list of the W matrix elements
     verbose : bool, optional
@@ -124,50 +131,58 @@ def cov_correct(hs, gxs, hdata, gdata, weights, Nxs, fvals, fcov, dw_dW_fs, Wval
 
     Notes
     -----
-    This function corresponds to the weighted score function with sweights (both terms of
-    Eq.51 in `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_). If the shape parameters
-    of the `gxs` are not known then the full sandwich estimate must be used which is not yet
-    implemented in this package.
+    This function corresponds to the weighted score function with sweights
+    (both terms of Eq.51 in
+    `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_).
+    If the shape parameters of the `gxs` are not known then the full sandwich
+    estimate must be used which is not yet implemented in this package.
 
     """
-
     dim_kl = len(fvals)
-    assert(fcov.shape[0]==dim_kl and fcov.shape[1]==dim_kl )
+    assert fcov.shape[0] == dim_kl and fcov.shape[1] == dim_kl
 
     dim_xy = len(gxs)
-    assert( len(Nxs)==dim_xy )
-    dim_E = int(dim_xy*(dim_xy+1)/2)  # indepdent elements of symmetric Wxy matrix is n(n+1)/2
-    assert( len(dw_dW_fs)==dim_E )
-    assert( len(Wvals)==dim_E )
+    assert len(Nxs) == dim_xy
+    dim_E = int(
+        dim_xy * (dim_xy + 1) / 2
+    )  # indepdent elements of symmetric Wxy matrix is n(n+1)/2
+    assert len(dw_dW_fs) == dim_E
+    assert len(Wvals) == dim_E
 
     HHpH_term = approx_cov_correct(hs, hdata, weights, fvals, fcov, verbose=False)
 
     # now construct the E and C' matrices
-    Ekl = np.empty((dim_kl,dim_E))
-    for l in range(dim_kl):
-        derivl = partial_derivative(hs,l,fvals,hdata)
-        gxevs = [ gx(gdata) for gx in gxs ]
+    Ekl = np.empty((dim_kl, dim_E))
+    for j in range(dim_kl):
+        derivj = _partial_derivative(hs, j, fvals, hdata)
+        gxevs = [gx(gdata) for gx in gxs]
         for xy in range(dim_E):
-            Ekl[l,xy] = np.sum( dw_dW_fs[xy]( *Wvals, *gxevs ) * derivl )
+            Ekl[j, xy] = np.sum(dw_dW_fs[xy](*Wvals, *gxevs) * derivj)
 
-    Ckl = np.empty((dim_E,dim_E))
-    gtot = np.sum( [ Nxs[i]*gxs[i](gdata) for i in range(dim_xy) ], axis=0 )
+    Ckl = np.empty((dim_E, dim_E))
+    gtot = np.sum([Nxs[i] * gxs[i](gdata) for i in range(dim_xy)], axis=0)
 
-    Citer = [ (gxs[i], gxs[j]) for i in range(dim_xy) for j in range(dim_xy) if i>=j ]
+    Citer = [(gxs[i], gxs[j]) for i in range(dim_xy) for j in range(dim_xy) if i >= j]
 
     for i, gis in enumerate(Citer):
         for j, gjs in enumerate(Citer):
-            Ckl[i,j] = np.sum( gis[0](gdata) * gis[1](gdata) * gjs[0](gdata) * gjs[1](gdata) / gtot**4 )
+            Ckl[i, j] = np.sum(
+                gis[0](gdata)
+                * gis[1](gdata)
+                * gjs[0](gdata)
+                * gjs[1](gdata)
+                / gtot ** 4
+            )
 
-    HECEH_term = fcov @ ( Ekl @ Ckl @ Ekl.T ) @ fcov.T
+    HECEH_term = fcov @ (Ekl @ Ckl @ Ekl.T) @ fcov.T
 
     tcov = HHpH_term - HECEH_term
 
     if verbose:
-        print('Full covariance correction for weighted events')
-        print('  Original covariance:')
-        print('\t', str(fcov).replace('\n','\n\t '))
-        print('  Corrected covariance:')
-        print('\t', str(tcov).replace('\n','\n\t '))
+        print("Full covariance correction for weighted events")
+        print("  Original covariance:")
+        print("\t", str(fcov).replace("\n", "\n\t "))
+        print("  Corrected covariance:")
+        print("\t", str(tcov).replace("\n", "\n\t "))
 
     return tcov
diff --git a/src/sweights/cow.py b/src/sweights/cow.py
index 0ff1470..7763a19 100644
--- a/src/sweights/cow.py
+++ b/src/sweights/cow.py
@@ -1,127 +1,138 @@
+"""Implementation of the COW class."""
+
 from scipy.stats import uniform
 from scipy.integrate import quad
 from scipy import linalg
 import numpy as np
 
-class cow():
+
+class cow:
     """Produce weights using COWs."""
 
     def __init__(self, mrange, gs, gb, Im=1, obs=None, renorm=True, verbose=True):
         """
         Initialize cow object.
 
-        This will compute the W and A (or alpha) matrices which are used to produce the weight
-        functions. Evaluation of these functions on a dataset is done in a different function
-        `getWeight`.
+        This will compute the W and A (or alpha) matrices which are used to
+        produce the weight functions. Evaluation of these functions on a
+        dataset is done in a different function `getWeight`.
 
         Parameters
         ----------
         mrange : tuple
-            A two element tuple for the integration range in the discriminant variable
+            A two element tuple for the integration range in the discriminant
+            variable
         gs : callable
-            The function for the signal pdf (numerator) must accept a single argument in
-            this case the discriminant variable
+            The function for the signal pdf (numerator) must accept a single
+            argument in this case the discriminant variable
         gb : callable or list of callable
-            A function or list of functions for the backgrond pdfs (numerator) which must
-            each accept a single argument in this case the discriminant variable
+            A function or list of functions for the backgrond pdfs (numerator)
+            which must each accept a single argument in this case the
+            discriminant variable
         Im : int or callable, optional
-            The function for the "variance function" or I(m) (denominator) which must accept
-            a single argument in this case the discriminant variable. Can also pass 1 for a
-            uniform variance function (the default)
+            The function for the "variance function" or I(m) (denominator)
+            which must accept a single argument in this case the discriminant
+            variable. Can also pass 1 for a uniform variance function (the
+            default)
         obs : tuple of ndarray, optional
-            You can instead pass the observed distribution to evaluate Im instead. This expects
-            the entries and bin edges in a two element tuple like the return value of np.histogram
+            You can instead pass the observed distribution to evaluate Im
+            instead. This expects the entries and bin edges in a two element
+            tuple like the return value of np.histogram
         renorm : bool, optional
-            Renormalise passed functions to unity (you can override this if you already know it's true)
+            Renormalise passed functions to unity (you can override this if you
+            already know it's true)
 
         Notes
         -----
-        For more details see `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_
+        For more details see
+        `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_
 
         See Also
         --------
         getWeight
         """
-
         self.renorm = renorm
         self.mrange = mrange
-        self.gs = self.normalise(gs)
-        self.gb = [ self.normalise(g) for g in gb ] if hasattr(gb,'__iter__') else [ self.normalise(gb) ]
+        self.gs = self._normalise(gs)
+        self.gb = (
+            [self._normalise(g) for g in gb]
+            if hasattr(gb, "__iter__")
+            else [self._normalise(gb)]
+        )
         self.gk = [self.gs] + self.gb
         if Im == 1:
             un = uniform(*mrange)
-            n  = np.diff( un.cdf(mrange) )
+            n = np.diff(un.cdf(mrange))
             self.Im = lambda m: un.pdf(m) / n
         else:
-            self.Im = self.normalise(Im)
+            self.Im = self._normalise(Im)
 
         self.obs = obs
         if obs:
-            if len(obs)!=2: raise ValueError('The observation must be passed as length two object containing weights and bin edges (w,xe) - ie. what is returned by np.histogram()')
+            if len(obs) != 2:
+                raise ValueError(
+                    """The observation must be passed as length two object
+                    containing weights and bin edges (w,xe) - ie. what is
+                    returned by np.histogram()"""
+                )
             w, xe = obs
-            if len(w)!=len(xe)-1: raise ValueError('The bin edges and weights do not have the right respective dimensions')
-# normalise
-            w = w/np.sum(w)  # sum of wts now 1
-            w /= (mrange[1]-mrange[0])/len(w) # now divide by bin width to get a function which integrates to 1
-            f = lambda m: w[ np.argmin( m >= xe )-1 ]
+            if len(w) != len(xe) - 1:
+                raise ValueError(
+                    """The bin edges and weights do not have the right
+                    respective dimensions"""
+                )
+            # normalise
+            w = w / np.sum(w)  # sum of wts now 1
+            w /= (mrange[1] - mrange[0]) / len(
+                w
+            )  # divide by bin width to get a function which integrates to 1
+
+            def f(m):
+                return w[np.argmin(m >= xe) - 1]
+
             self.Im = np.vectorize(f)
 
-        if verbose: print('Initialising COW:')
+        if verbose:
+            print("Initialising COW:")
 
         # compute Wkl matrix
-        self.Wkl = self.compWkl()
+        self.Wkl = self._compWkl()
         if verbose:
-            print('    W-matrix:')
-            print( '\t'+ str(self.Wkl).replace('\n','\n\t ') )
+            print("    W-matrix:")
+            print("\t" + str(self.Wkl).replace("\n", "\n\t "))
 
         # invert for Akl matrix
-        self.Akl = linalg.solve( self.Wkl, np.identity( len(self.Wkl) ), assume_a='pos' )
+        self.Akl = linalg.solve(self.Wkl, np.identity(len(self.Wkl)), assume_a="pos")
         if verbose:
-            print('    A-matrix:')
-            print( '\t'+ str(self.Akl).replace('\n','\n\t ') )
-
-    def normalise(self, f):
-        """
-        Return a normalised pdf
-
-        Parameters
-        ----------
-        f : callable
-            Input function
+            print("    A-matrix:")
+            print("\t" + str(self.Akl).replace("\n", "\n\t "))
 
-        Returns
-        -------
-        callable :
-            Normalised output function
-        """
+    def _normalise(self, f):
         if self.renorm:
-            N = quad(f,*self.mrange)[0]
+            N = quad(f, *self.mrange)[0]
             return lambda m: f(m) / N
         else:
             return f
 
-    def compWklElem(self, k, l):
+    def _compWklElem(self, k, j):
 
         # check it's available in m
-        assert( k < len(self.gk) )
-        assert( l < len(self.gk) )
+        assert k < len(self.gk)
+        assert j < len(self.gk)
 
         def integral(m):
-            return self.gk[k](m) * self.gk[l](m) / self.Im(m)
+            return self.gk[k](m) * self.gk[j](m) / self.Im(m)
 
         if self.obs is None:
-            return quad( integral, self.mrange[0], self.mrange[1] )[0]
+            return quad(integral, self.mrange[0], self.mrange[1])[0]
         else:
             tint = 0
             xe = self.obs[1]
-            for le, he in zip(xe[:-1],xe[1:]):
-                tint += quad( integral, le, he )[0]
+            for le, he in zip(xe[:-1], xe[1:]):
+                tint += quad(integral, le, he)[0]
             return tint
 
-    def compWkl(self):
-        """
-        Compute the W matrix
-        """
+    def _compWkl(self):
 
         n = len(self.gk)
 
@@ -129,14 +140,16 @@ def compWkl(self):
 
         for i in range(n):
             for j in range(n):
-                if i>j: ret[i,j] = ret[j,i]
-                else:   ret[i,j] = self.compWklElem(i,j)
+                if i > j:
+                    ret[i, j] = ret[j, i]
+                else:
+                    ret[i, j] = self._compWklElem(i, j)
 
         return ret
 
     def wk(self, k, m):
         """
-        Return the weights
+        Return the weights.
 
         Parameters
         ----------
@@ -150,13 +163,14 @@ def wk(self, k, m):
         ndarray :
             Values of the weights
         """
-
         n = len(self.gk)
-        return np.sum( [ self.Akl[k,l] * self.gk[l](m) / self.Im(m) for l in range(n) ], axis=0 )
+        return np.sum(
+            [self.Akl[k, j] * self.gk[j](m) / self.Im(m) for j in range(n)], axis=0
+        )
 
     def getWeight(self, k, m):
         """
-        Return the weights
+        Return the weights.
 
         Wrapper for `wk`
 
@@ -164,7 +178,4 @@ def getWeight(self, k, m):
         --------
         wk
         """
-        return self.wk(k,m)
-
-
-
+        return self.wk(k, m)
diff --git a/src/sweights/independence.py b/src/sweights/independence.py
index 2441329..dc3ef28 100644
--- a/src/sweights/independence.py
+++ b/src/sweights/independence.py
@@ -1,7 +1,9 @@
+"""Module to check and plot independence."""
 import numpy as np
 from scipy.stats import kendalltau
 
-def kendall_tau(x,y):
+
+def kendall_tau(x, y):
     """
     Return kendall tau correlation coefficient.
 
@@ -30,44 +32,65 @@ def kendall_tau(x,y):
     calcualtion (the uncertainty calculation is trivial) which makes a
     few optimisations. See the scipy documentation for more information.
     """
-    assert(len(x)==len(y))
-    err_approx = 1./np.sqrt(len(x))
-    return ( kendalltau(x,y).correlation, err_approx, kendalltau(x,y).pvalue )
+    assert len(x) == len(y)
+    err_approx = 1.0 / np.sqrt(len(x))
+    return (kendalltau(x, y).correlation, err_approx, kendalltau(x, y).pvalue)
 
-    raise RuntimeWarning('This function is depreciated use scipy.stats.kendalltau instead')
-    factor = 2./(len(x)*(len(x)-1))
-    su = 0.
+    raise RuntimeWarning(
+        "This function is depreciated use scipy.stats.kendalltau instead"
+    )
+    factor = 2.0 / (len(x) * (len(x) - 1))
+    su = 0.0
     for i in range(len(x)):
-        for j in range(i,len(x)):
-            su += np.sign(x[i]-x[j])*np.sign(y[i]-y[j])
+        for j in range(i, len(x)):
+            su += np.sign(x[i] - x[j]) * np.sign(y[i] - y[j])
+
+    return (factor * su, err_approx)
 
-    return (factor*su, err_approx)
 
-def plot(x,y,save=None,show=False):
+def plot(x, y, save=None, show=False):
+    """
+    Plot scatter of two variables.
+
+    Plot a scatter graph of two variables and write the kendall tau
+    coefficient.
+    """
     try:
-      import matplotlib.pyplot as plt
-      import uncertainties as u
-    except:
-      raise RuntimeError('matplotlib and uncertainties packages must be installed to plot independence \npip install matplotlib \npip install uncertainties')
+        import matplotlib.pyplot as plt
+        import uncertainties as u
+    except Exception:
+        raise RuntimeError(
+            """matplotlib and uncertainties packages must be installed to plot
+            independence \npip install matplotlib \npip install
+            uncertainties"""
+        )
 
     fig, ax = plt.subplots()
-    ax.scatter(x,y)
-    tau, err = kendall_tau(x,y)
-    ax.text(0.7,0.9, r'$\tau = '+str(u.ufloat(tau,err)).replace('+/-',r'\pm')+'$', transform=ax.transAxes, backgroundcolor='w')
-    ax.set_xlabel('x')
-    ax.set_ylabel('y')
+    ax.scatter(x, y)
+    tau, err = kendall_tau(x, y)
+    ax.text(
+        0.7,
+        0.9,
+        r"$\tau = " + str(u.ufloat(tau, err)).replace("+/-", r"\pm") + "$",
+        transform=ax.transAxes,
+        backgroundcolor="w",
+    )
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
     fig.tight_layout()
-    if save: fig.savefig(save)
-    if show: plt.show()
+    if save:
+        fig.savefig(save)
+    if show:
+        plt.show()
     return fig, ax
 
-if __name__=="__main__":
+
+if __name__ == "__main__":
 
     a = list(range(10))
     b = list(reversed(range(10)))
-    #plot(a,b)
+    # plot(a,b)
 
     x = np.random.uniform(size=1000)
     y = np.random.uniform(size=1000)
-    plot(x,y)
-
+    plot(x, y)
diff --git a/src/sweights/sweight.py b/src/sweights/sweight.py
index 044774e..6a586d7 100644
--- a/src/sweights/sweight.py
+++ b/src/sweights/sweight.py
@@ -1,256 +1,331 @@
-# vim: ts=4 sw=4
+"""Implementation of the sweight class."""
 
 import numpy as np
 from scipy.integrate import nquad
 from scipy.linalg import solve
 from scipy.interpolate import InterpolatedUnivariateSpline
 
-class sweight():
+
+class sweight:
     """Produce sweights for a dataset given component pdfs."""
 
-    def __init__(self,
-                 data,
-                 pdfs=[],
-                 yields=[],
-                 discvarranges=None,
-                 method='summation',
-                 compnames=None,
-                 alphas=None,
-                 rfobs=[],
-                 verbose=True,
-                 checks=True
-                ):
+    def __init__(
+        self,
+        data,
+        pdfs=[],
+        yields=[],
+        discvarranges=None,
+        method="summation",
+        compnames=None,
+        alphas=None,
+        rfobs=[],
+        verbose=True,
+        checks=True,
+    ):
         """
         Initialize sweight object.
 
-        This will compute the W and A (or alpha) matrices which are used to produce the weight
-        functions. Evaluation of these functions on a dataset is done in a different function
-        `getWeight`.
+        This will compute the W and A (or alpha) matrices which are used to
+        produce the weight functions. Evaluation of these functions on a
+        dataset is done in a different function `getWeight`.
 
         Parameters
         ----------
         data : ndarray
-            The dataset in the discriminating variable, should have shape (nevents,ndiscvars)
-            so for one discriminating variable will just be (N,) or (N,1).
+            The dataset in the discriminating variable, should have shape
+            (nevents,ndiscvars) so for one discriminating variable will just
+            be (N,) or (N,1).
         pdfs : list of callable
-            A list of the component pdfs. These should be simple python callable functions which
-            are passed the same number of parameters as there are discriminating variables.
-            If running with the 'roofit' method (see `method`) then this can be a list of
-            ROOT.RooAbsPdf objects. If you only have your pdfs defined as RooFit objects then you
-            can use the wrapper function `convertRooAbsPdf` to convert them into the appropriate
-            object type for this call and then use other methods.
+            A list of the component pdfs. These should be simple python
+            callable functions which are passed the same number of parameters
+            as there are discriminating variables. If running with the
+            'roofit' method (see `method`) then this can be a list of
+            ROOT.RooAbsPdf objects. If you only have your pdfs defined as
+            RooFit objects then you can use the wrapper function
+            `convertRooAbsPdf` to convert them into the appropriate object type
+            for this call and then use other methods.
         yields : list of float
             A list of the component yields.
         discvarranges : list of tuple or tuple of tuple, optional
-            The ranges for each discriminating variable, as a tuple or list of two element tuples
-            specifying the lower and upper bound for the range (the default is None which will use
-            minus to plus infinity)
+            The ranges for each discriminating variable, as a tuple or list of
+            two element tuples specifying the lower and upper bound for the
+            range (the default is None which will use minus to plus infinity)
         method: str, optional
-            The sweights method to use. Must be one of 'summation', 'integration', 'subhess', 'tsplot',
-            'rootfit'. The recommended default is summation corresponding to Variant B of
+            The sweights method to use. Must be one of 'summation',
+            'integration', 'subhess', 'tsplot', 'rootfit'. The recommended
+            default is summation corresponding to Variant B of
             `arXiv:2112.04574 <https://arxiv.org/abs/2112.04574>`_
         compnames: list of str, optional
-            A list of the component names. Only used for the legend entries when making a plot of the
-            weight functions with `makeWeightPlot`
+            A list of the component names. Only used for the legend entries
+            when making a plot of the weight functions with `makeWeightPlot`
         alphas: ndarray, optional
-            If using the 'subhess' method then the covariance matrix of a fit to the disciminanting
-            variable(s) in which only the yields float must also be passed. This matrix is inverted
-            to produce the W-matrix.
+            If using the 'subhess' method then the covariance matrix of a fit
+            to the disciminanting variable(s) in which only the yields float
+            must also be passed. This matrix is inverted to produce the
+            W-matrix.
         rfobs: list, optional
-            If using the 'roofit' method then the discriminating variables (of type RooRealVar) on which
-            the pdfs depend must also be passed
+            If using the 'roofit' method then the discriminating variables
+            (of type RooRealVar) on which the pdfs depend must also be passed
         verbose: bool, optional
             Print output
         checks: bool, optional
-            Perform some checks that the derived weights are self consistent, in particular check that
-            the sum of all component weights for a given value of the discriminating variable(s) is unity,
-            and check that the sum of all weights for a given component reproduces the yields. This will
-            print additional output to the screen and issue warnings if checks are not passed.
+            Perform some checks that the derived weights are self consistent,
+            in particular check that the sum of all component weights for a
+            given value of the discriminating variable(s) is unity, and check
+            that the sum of all weights for a given component reproduces the
+            yields. This will print additional output to the screen and issue
+            warnings if checks are not passed.
 
         Notes
         -----
-        Many of the arguments passed and methods implemented are not strictly necessary. In particular the
-        'tsplot' and 'roofit' methods are just wrappers for implementations elsewhere that were originally
-        used as cross-checks. In future versions these two methods should be removed to simplify this call.
+        Many of the arguments passed and methods implemented are not strictly
+        necessary. In particular the 'tsplot' and 'roofit' methods are just
+        wrappers for implementations elsewhere that were originally used as
+        cross-checks. In future versions these two methods should be removed
+        to simplify this call.
 
         See Also
         --------
         getWeight, makeWeightPlot
         """
-
-        self.allowed_methods = ['summation','integration','refit','subhess','tsplot','roofit']
+        self.allowed_methods = [
+            "summation",
+            "integration",
+            "refit",
+            "subhess",
+            "tsplot",
+            "roofit",
+        ]
         self.method = method
         if self.method not in self.allowed_methods:
-            raise RuntimeError( self.method,'is not an allowed method. Must be one of', self.allowed_methods )
-
-        if self.method=='refit':
-          try:
-            import iminuit
-            from iminuit import Minuit
-          except:
-            raise RuntimeError( 'To run sweight with refit method iminuit must be installed')
-          if iminuit.version.version.split('.')[0]=='1':
-            raise RuntimeError( 'iminuit version must be > 2.x' )
-
-        if self.method=='tsplot':
-          try:
-            import ROOT as r
-          except:
-            raise RuntimeError( 'To run sweight with tsplot method ROOT must be installed')
-
-        if self.method=='roofit':
-          try:
-            import ROOT as r
-            from ROOT import RooFit as rf
-            from ROOT import RooStats as rs
-          except:
-            raise RuntimeError( 'To run sweight with roofit method ROOT with RooFit must be installed')
-
+            raise RuntimeError(
+                self.method,
+                "is not an allowed method. Must be one of",
+                self.allowed_methods,
+            )
+
+        if self.method == "refit":
+            try:
+                import iminuit
+            except Exception:
+                raise RuntimeError(
+                    "To run with the 'refit' method iminuit must be installed"
+                )
+            if iminuit.version.version.split(".")[0] == "1":
+                raise RuntimeError("iminuit version must be > 2.x")
+
+        if self.method == "tsplot":
+            try:
+                import ROOT as r
+            except Exception:
+                raise RuntimeError("To run with 'tsplot method ROOT must be installed")
+            if int(r.__version__.split(".")[0]) < 6:
+                raise RuntimeError("ROOT version must be > 6.xx")
+
+        if self.method == "roofit":
+            try:
+                import ROOT as r
+            except Exception:
+                raise RuntimeError(
+                    """To run sweight with roofit method ROOT with RooFit must
+                    be installed"""
+                )
+            if int(r.__version__.split(".")[0]) < 6:
+                raise RuntimeError("ROOT version must be > 6.xx")
 
         self.verbose = verbose
-        if self.verbose: print('Initialising sweight with the', self.method, 'method:')
+        if self.verbose:
+            print("Initialising sweight with the", self.method, "method:")
 
         # get data in right format
         self.data = data
-        if len(self.data.shape)==1:
-            self.data = self.data.reshape( len(self.data), 1 )
-        if not len(self.data.shape)==2:
-            raise ValueError( 'Input data is in the wrong format. Should be a numpy ndarray with shape (nevs,ncomps)' )
+        if len(self.data.shape) == 1:
+            self.data = self.data.reshape(len(self.data), 1)
+        if not len(self.data.shape) == 2:
+            raise ValueError(
+                """Input data is in the wrong format. Should be a numpy
+                ndarray with shape (nevs,ncomps)"""
+            )
         self.nevs = self.data.shape[0]
         self.ndiscvars = self.data.shape[1]
 
         # sort out the disciminant variables integration ranges
         self.discvarranges = discvarranges
         if self.discvarranges is None:
-            self.discvarranges = tuple( tuple(-np.inf,np.inf) for i in range(self.ndiscvars) )
-        if not (self.ndiscvars==len(self.discvarranges)):
-            raise ValueError( 'You dont seemed to have passed sufficient ranges', len(self.discvarranges),'for the number of discriminant variables', len(self.ndiscvars))
-        self.discvarranges = np.array( self.discvarranges )
+            self.discvarranges = tuple(
+                tuple(-np.inf, np.inf) for i in range(self.ndiscvars)
+            )
+        if not (self.ndiscvars == len(self.discvarranges)):
+            raise ValueError(
+                "You dont seemed to have passed sufficient ranges",
+                len(self.discvarranges),
+                "for the number of discriminant variables",
+                len(self.ndiscvars),
+            )
+        self.discvarranges = np.array(self.discvarranges)
 
         # remove events out of range
-        self.data = self.data[ np.logical_and( self.data >= self.discvarranges.T[0], self.data <= self.discvarranges.T[1] ).all( axis=1 ) ]
+        self.data = self.data[
+            np.logical_and(
+                self.data >= self.discvarranges.T[0],
+                self.data <= self.discvarranges.T[1],
+            ).all(axis=1)
+        ]
 
         # check pdfs and write normalisation terms
         self.pdfs = pdfs
         self.yields = yields
-        if not (len(self.yields)==len(self.pdfs)):
-          raise ValueError( 'The number of yields is not the same as the number of pdfs' )
+        if not (len(self.yields) == len(self.pdfs)):
+            raise ValueError(
+                "The number of yields is not the same as the number of pdfs"
+            )
 
         self.ncomps = len(self.yields)
-        if self.method != 'roofit':
-          self.pdfnorms = [ nquad( pdf, self.discvarranges )[0] for pdf in self.pdfs ]
-          if checks:
-            print('    PDF normalisations:')
-            for i, norm in enumerate(self.pdfnorms):
-              print('\t', i, norm)
-          self.pdfsum   = lambda x: sum( [self.yields[i]*self.pdfs[i](x)/self.pdfnorms[i] for i in range(self.ncomps)] )
+        if self.method != "roofit":
+            self.pdfnorms = [nquad(pdf, self.discvarranges)[0] for pdf in self.pdfs]
+            if checks:
+                print("    PDF normalisations:")
+                for i, norm in enumerate(self.pdfnorms):
+                    print("\t", i, norm)
+            self.pdfsum = lambda x: sum(
+                [
+                    self.yields[i] * self.pdfs[i](x) / self.pdfnorms[i]
+                    for i in range(self.ncomps)
+                ]
+            )
 
         # if subhess method check alpha has been passed
         self.alphas = alphas
-        if self.method=='subhess':
+        if self.method == "subhess":
             if alphas is None:
-                raise RuntimeError('If using method subhess you must pass the covariance (alpha) matrix')
-        if alphas is not None and not alphas.shape==np.zeros( (self.ncomps,self.ncomps) ).shape:
-            raise RuntimeError('You have passed an alpha matrix but it has the wrong shape')
+                raise RuntimeError(
+                    """If using method subhess you must pass the covariance
+                    (alpha) matrix"""
+                )
+        if (
+            alphas is not None
+            and not alphas.shape == np.zeros((self.ncomps, self.ncomps)).shape
+        ):
+            raise RuntimeError(
+                "You have passed an alpha matrix but it has the wrong shape"
+            )
 
         # add names if needed
         self.compnames = compnames
         if self.compnames is None:
-          self.compnames = [ str(i) for i in range(self.ncomps) ]
+            self.compnames = [str(i) for i in range(self.ncomps)]
 
         # compute the W matrix
-        self.computeWMatrix()
+        self._computeWMatrix()
 
         # solve the alpha matrix
-        self.solveAlphas()
+        self._solveAlphas()
 
         # do the tsplot implementation if asked
         self.tsplotweights = None
         self.tsplotw = None
-        if self.method=='tsplot':
-          self.tsplotweights = self.runTSPlot()
-          self.tsplotw = [ InterpolatedUnivariateSpline( *self.tsplotweights[:,[0,i+1]].T, k=3 ) for i in range(self.ncomps) ]
+        if self.method == "tsplot":
+            self.tsplotweights = self._runTSPlot()
+            self.tsplotw = [
+                InterpolatedUnivariateSpline(*self.tsplotweights[:, [0, i + 1]].T, k=3)
+                for i in range(self.ncomps)
+            ]
 
         # do the RooFit implementation if asked
         self.roofitweights = None
         self.roofitw = None
         self.rfobs = rfobs
-        if self.method=='roofit':
-          if len(self.rfobs) != self.ndiscvars:
-              raise RuntimeError('You must pass a list of RooFit observables the same length as the data when running with the roofit method')
-
-          self.roofitweights = self.runRooFit()
-          self.roofitw = [ InterpolatedUnivariateSpline( *self.roofitweights[:,[0,i+1]].T, k=3 ) for i in range(self.ncomps) ]
+        if self.method == "roofit":
+            if len(self.rfobs) != self.ndiscvars:
+                raise RuntimeError(
+                    """You must pass a list of RooFit observables the same
+                    length as the data when running with the roofit method"""
+                )
+
+            self.roofitweights = self._runRooFit()
+            self.roofitw = [
+                InterpolatedUnivariateSpline(*self.roofitweights[:, [0, i + 1]].T, k=3)
+                for i in range(self.ncomps)
+            ]
 
         # print checks
-        if checks: self.printChecks()
+        if checks:
+            self.printChecks()
 
-    def computeWMatrix(self):
-        """
-        Compute the W matrix
-        """
-        self.Wkl = np.zeros( (self.ncomps,self.ncomps) )
-        if self.method in ['refit','subhess','tsplot','roofit']:
-          return self.Wkl
+    def _computeWMatrix(self):
+        self.Wkl = np.zeros((self.ncomps, self.ncomps))
+        if self.method in ["refit", "subhess", "tsplot", "roofit"]:
+            return self.Wkl
         for i, di in enumerate(self.pdfs):
             dinorm = self.pdfnorms[i]
             for j, dj in enumerate(self.pdfs):
                 djnorm = self.pdfnorms[j]
-                if j<i:
-                    self.Wkl[i,j] = self.Wkl[j,i]
+                if j < i:
+                    self.Wkl[i, j] = self.Wkl[j, i]
                 else:
-                    if self.method=='integration':
-                        self.Wkl[i,j] = nquad( lambda *args: ( di(*args)/dinorm  * dj(*args)/djnorm ) / self.pdfsum(*args), self.discvarranges )[0]
-                    elif self.method=='summation':
-                        self.Wkl[i,j] = np.sum( ( di(*self.data.T)/dinorm * dj(*self.data.T)/djnorm ) / self.pdfsum(*self.data.T)**2 )
+                    if self.method == "integration":
+                        self.Wkl[i, j] = nquad(
+                            lambda *args: (di(*args) / dinorm * dj(*args) / djnorm)
+                            / self.pdfsum(*args),
+                            self.discvarranges,
+                        )[0]
+                    elif self.method == "summation":
+                        self.Wkl[i, j] = np.sum(
+                            (di(*self.data.T) / dinorm * dj(*self.data.T) / djnorm)
+                            / self.pdfsum(*self.data.T) ** 2
+                        )
                     else:
-                        self.Wkl[i,j] = 0.
+                        self.Wkl[i, j] = 0.0
 
-        if self.method in ['integration','summation'] and self.verbose:
-            print('    W-matrix:')
-            print('\t'+ str(self.Wkl).replace('\n','\n\t'))
+        if self.method in ["integration", "summation"] and self.verbose:
+            print("    W-matrix:")
+            print("\t" + str(self.Wkl).replace("\n", "\n\t"))
 
         return self.Wkl
 
-    def nll(self, pars):
-        """
-        Define NLL function to minimise if 'refit' method is used
-        """
-        assert(len(pars)==self.ncomps)
+    def _nll(self, pars):
+        assert len(pars) == self.ncomps
         nobs = sum(pars)
-        nest = np.sum( np.log( sum( [ pars[i]*pdf(*self.data.T)/self.pdfnorms[i] for i, pdf in enumerate(self.pdfs) ] ) ) )
+        nest = np.sum(
+            np.log(
+                sum(
+                    [
+                        pars[i] * pdf(*self.data.T) / self.pdfnorms[i]
+                        for i, pdf in enumerate(self.pdfs)
+                    ]
+                )
+            )
+        )
         return nobs - nest
 
-    def solveAlphas(self):
-        """
-        Compute the A (or alpha) matrix
+    def _solveAlphas(self):
+        if self.method in ["integration", "summation"]:
+            sol = np.identity(len(self.Wkl))
+            self.alphas = solve(self.Wkl, sol, assume_a="pos")
+        elif self.method in ["refit"]:
+            from iminuit import Minuit
 
-        Normally this is done by simply inverting the Wkl matrix
-        """
-        if self.method in ['integration','summation']:
-            sol = np.identity( len(self.Wkl) )
-            self.alphas = solve( self.Wkl, sol, assume_a='pos' )
-        elif self.method in ['refit']:
-            #mi = Minuit.from_array_func(self.nll, tuple(self.yields), errordef=Minuit.LIKELIHOOD, pedantic=False)
-            mi = Minuit(self.nll, tuple(self.yields) )
-            mi.errordef=Minuit.LIKELIHOOD
+            mi = Minuit(self._nll, tuple(self.yields))
+            mi.errordef = Minuit.LIKELIHOOD
             mi.migrad()
             mi.hesse()
             self.alphas = np.array(mi.covariance.tolist())
-        elif self.method in ['tsplot','roofit']:
-            self.alphas = np.identity( len(self.Wkl) )
+        elif self.method in ["tsplot", "roofit"]:
+            self.alphas = np.identity(len(self.Wkl))
             return self.alphas
 
         if self.verbose:
-            print('    A-matrix:')
-            print('\t' + str(self.alphas).replace('\n','\n\t'))
+            print("    A-matrix:")
+            print("\t" + str(self.alphas).replace("\n", "\n\t"))
 
         return self.alphas
 
-    def getWeight(self,icomp=0, *args):
+    def getWeight(self, icomp=0, *args):
         """
+        Return the weights.
+
         Get the weights for a given component and set of discriminating
-        variable values
+        variable values.
 
         Parameters
         ----------
@@ -265,16 +340,28 @@ def getWeight(self,icomp=0, *args):
         ndarray
             An array of the weights
         """
-        if self.method == 'tsplot':
+        if self.method == "tsplot":
             return self.tsplotw[icomp](*args)
-        elif self.method == 'roofit':
+        elif self.method == "roofit":
             return self.roofitw[icomp](*args)
         else:
-            return sum( [self.alphas[i,icomp] * self.pdfs[i](*args)/self.pdfnorms[i] for i in range(self.ncomps)] ) / sum( [self.yields[i] * self.pdfs[i](*args)/self.pdfnorms[i] for i in range(self.ncomps)] )
-
-    def makeWeightPlot(self, axis=None, dopts=['r','b','g','m','c','y'], labels=None):
+            return sum(
+                [
+                    self.alphas[i, icomp] * self.pdfs[i](*args) / self.pdfnorms[i]
+                    for i in range(self.ncomps)
+                ]
+            ) / sum(
+                [
+                    self.yields[i] * self.pdfs[i](*args) / self.pdfnorms[i]
+                    for i in range(self.ncomps)
+                ]
+            )
+
+    def makeWeightPlot(
+        self, axis=None, dopts=["r", "b", "g", "m", "c", "y"], labels=None
+    ):
         """
-        Make a plot of the weight functions
+        Make a plot of the weight functions.
 
         Parameters
         ----------
@@ -283,249 +370,303 @@ def makeWeightPlot(self, axis=None, dopts=['r','b','g','m','c','y'], labels=None
         dopts : list of str
             List of colors for the different components
         labels : list of str
-            List of legend labels. Default will use those passed in `compnames` with `__init__`
+            List of legend labels. Default will use those passed in `compnames`
+            with `__init__`
         """
-        if self.ndiscvars!=1:
-            print('WARNING - I dont know how to plot this')
+        if self.ndiscvars != 1:
+            print("WARNING - I dont know how to plot this")
             return None
 
-        while len(dopts)<self.ncomps:
-          dopts.append('b')
+        while len(dopts) < self.ncomps:
+            dopts.append("b")
 
         try:
-          import matplotlib.pyplot as plt
-        except:
-          raise RuntimeError('matplotlib must be installed to make the weight plot')
+            import matplotlib.pyplot as plt
+        except Exception:
+            raise RuntimeError("matplotlib must be installed to make the weight plot")
 
         ax = axis or plt.gca()
 
-        x = np.linspace( *self.discvarranges[0], 400 )
+        x = np.linspace(*self.discvarranges[0], 400)
 
-        labels = labels or [ '$w_{{{0}}}$'.format(comp) for comp in self.compnames ]
+        labels = labels or ["$w_{{{0}}}$".format(comp) for comp in self.compnames]
 
         for comp in range(self.ncomps):
-            ax.plot( x, self.getWeight(comp,x), color=dopts[comp], linewidth=2, label=labels[comp] )
-
-        label = labels[-1] if len(labels)>self.ncomps else '$\sum_i w_{i}$'
-        ax.plot( x, sum( [ self.getWeight(c,x) for c in range(self.ncomps) ] ), 'k-', linewidth=3, label=label)
+            ax.plot(
+                x,
+                self.getWeight(comp, x),
+                color=dopts[comp],
+                linewidth=2,
+                label=labels[comp],
+            )
+
+        label = labels[-1] if len(labels) > self.ncomps else r"$\sum_i w_{i}$"
+        ax.plot(
+            x,
+            sum([self.getWeight(c, x) for c in range(self.ncomps)]),
+            "k-",
+            linewidth=3,
+            label=label,
+        )
 
         ax.legend()
 
     def printChecks(self):
-        """
-        Print checks
-        """
-        if self.method!='roofit':
-            self.intws = np.identity( self.ncomps )
+        """Print checks."""
+        if self.method != "roofit":
+            self.intws = np.identity(self.ncomps)
             for i in range(self.ncomps):
                 for j in range(self.ncomps):
-                    self.intws[i,j] = nquad( lambda *args:  self.getWeight(i,*args) * self.pdfs[j](*args) / self.pdfnorms[j], self.discvarranges )[0]
-            print('    Integral of w*pdf matrix (should be close to the identity):')
-            #with np.printoptions(precision=3, suppress=True):
-            print( '\t' + str(self.intws).replace('\n','\n\t') )
-
-        print('    Check of weight sums (should match yields):')
-        self.sows = [ np.sum( self.getWeight(i,*self.data.T) ) for i in range(self.ncomps) ]
-        header = '\t{:10s} | {:^10s} | {:^10s} | {:^9s} |'.format('Component','sWeightSum','Yield','Diff')
+                    self.intws[i, j] = nquad(
+                        lambda *args: self.getWeight(i, *args)
+                        * self.pdfs[j](*args)
+                        / self.pdfnorms[j],
+                        self.discvarranges,
+                    )[0]
+            print(
+                """    Integral of w*pdf matrix (should be close to the
+                identity):"""
+            )
+            # with np.printoptions(precision=3, suppress=True):
+            print("\t" + str(self.intws).replace("\n", "\n\t"))
+
+        print("    Check of weight sums (should match yields):")
+        self.sows = [
+            np.sum(self.getWeight(i, *self.data.T)) for i in range(self.ncomps)
+        ]
+        header = "\t{:10s} | {:^10s} | {:^10s} | {:^9s} |".format(
+            "Component", "sWeightSum", "Yield", "Diff"
+        )
         print(header)
-        print( '\t'+'-'.join( ['' for i in range(len(header)+1)]) )
+        print("\t" + "-".join(["" for i in range(len(header) + 1)]))
 
         for i, yld in enumerate(self.yields):
-            print('\t  {:<8d} | {:10.4f} | {:10.4f} | {:8.2f}% |'.format(i, self.sows[i], yld, 100.*(yld-self.sows[i])/self.sows[i]))
+            print(
+                "\t  {:<8d} | {:10.4f} | {:10.4f} | {:8.2f}% |".format(
+                    i, self.sows[i], yld, 100.0 * (yld - self.sows[i]) / self.sows[i]
+                )
+            )
 
+    def _runTSPlot(self):
 
-    def runTSPlot(self):
-        """
-        Run the TSPlot specific method
-        """
+        import ROOT as r
 
-        if self.method!='tsplot':
-            raise RuntimeError('Method is',self.method,'but calling the tsplot function.')
+        if self.method != "tsplot":
+            raise RuntimeError(
+                "Method is", self.method, "but calling the tsplot function."
+            )
 
         # this works very differently for nD fits
         # so will not implement it here
-        if self.ndiscvars!=1:
-            raise RuntimeError('Sorry but I can\'t do the tsplot method for >1D fits')
+        if self.ndiscvars != 1:
+            raise RuntimeError("Sorry but I can't do the tsplot method for >1D fits")
 
         # make the tree for TSPlot
 
         # temporary datfile
-        datfile = open('.data.dat','w')
+        datfile = open(".data.dat", "w")
         for x in self.data:
-          datfile.write('{}'.format(x[0]))
-          for i, pdf in enumerate(self.pdfs):
-            x_pdf = pdf(x[0])/self.pdfnorms[i]
-            datfile.write(' {}'.format(x_pdf))
-          datfile.write('\n')
+            datfile.write("{}".format(x[0]))
+            for i, pdf in enumerate(self.pdfs):
+                x_pdf = pdf(x[0]) / self.pdfnorms[i]
+                datfile.write(" {}".format(x_pdf))
+            datfile.write("\n")
         datfile.close()
 
         # now make the tree and draw some plots to check it
         r.gROOT.SetBatch()
 
-        tree = r.TTree('data','data')
-        read_str = "x/D:" + ":".join( ["fx_%s"%a for a in self.compnames] )
+        tree = r.TTree("data", "data")
+        read_str = "x/D:" + ":".join(["fx_%s" % a for a in self.compnames])
 
-        tree.ReadFile('.data.dat', read_str, ' ')
-        c = r.TCanvas("c","c",(self.ncomps+1)*600,400)
-        c.Divide(self.ncomps+1,1)
+        tree.ReadFile(".data.dat", read_str, " ")
+        c = r.TCanvas("c", "c", (self.ncomps + 1) * 600, 400)
+        c.Divide(self.ncomps + 1, 1)
         c.cd(1)
         tree.Draw("x")
         for i in range(self.ncomps):
-          c.cd(i+2)
-          tree.Draw("fx_%s:x"%self.compnames[i])
+            c.cd(i + 2)
+            tree.Draw("fx_%s:x" % self.compnames[i])
         c.Update()
         c.Modified()
         c.Draw()
         r.gErrorIgnoreLevel = r.kInfo
-        c.Print('figs/tree.pdf')
+        c.Print("figs/tree.pdf")
 
         # now do the TSPlot
         from array import array
+
         tsplot = r.TSPlot(0, self.ndiscvars, len(self.data), self.ncomps, tree)
-        sel_str = "x:" + ":".join( ["fx_%s"%a for a in self.compnames] )
+        sel_str = "x:" + ":".join(["fx_%s" % a for a in self.compnames])
         tsplot.SetTreeSelection(sel_str)
-        ne = array('i', [int(yld) for yld in self.yields] )
+        ne = array("i", [int(yld) for yld in self.yields])
         tsplot.SetInitialNumbersOfSpecies(ne)
         tsplot.MakeSPlot("Q")
 
         # get the weights out of TSPlot
-        weights = np.ndarray( len(self.data)*self.ncomps )
+        weights = np.ndarray(len(self.data) * self.ncomps)
         tsplot.GetSWeights(weights)
-        weights = np.reshape( weights, (-1, self.ncomps) )
-        data_w_weights = np.append( self.data, weights, axis=1 )
-        sorted_data_w_weights = data_w_weights[np.argsort(data_w_weights[:,0])]
+        weights = np.reshape(weights, (-1, self.ncomps))
+        data_w_weights = np.append(self.data, weights, axis=1)
+        sorted_data_w_weights = data_w_weights[np.argsort(data_w_weights[:, 0])]
 
         return sorted_data_w_weights
 
-    def runRooFit(self):
-        """
-        Run the RooFit specific method
-        """
+    def _runRooFit(self):
 
-        if self.method!='roofit':
-            raise RuntimeError('Method is',self.method,'but calling the roofit function.')
+        import ROOT as r
 
-        r.RooMsgService.instance().setGlobalKillBelow( rf.FATAL )
+        if self.method != "roofit":
+            raise RuntimeError(
+                "Method is", self.method, "but calling the roofit function."
+            )
+
+        r.RooMsgService.instance().setGlobalKillBelow(r.RooFit.FATAL)
         r.gErrorIgnoreLevel = r.kInfo
 
         # sort out observables
-        rf_obs  = r.RooArgList()
+        rf_obs = r.RooArgList()
         for obs in self.rfobs:
-            if not obs.InheritsFrom('RooAbsReal'):
-                raise RuntimeError('Found an observable which does not inherit from RooAbsReal')
-            rf_obs.add( obs )
+            if not obs.InheritsFrom("RooAbsReal"):
+                raise RuntimeError(
+                    """Found an observable which does not inherit from
+                    RooAbsReal"""
+                )
+            rf_obs.add(obs)
 
         # make the roodataset and fill it
-        rf_dset = r.RooDataSet( 'data', 'data', r.RooArgSet(rf_obs) )
+        rf_dset = r.RooDataSet("data", "data", r.RooArgSet(rf_obs))
         for row in self.data:
-          for i in range(self.ndiscvars):
-            rf_obs.at(i).setVal( row[i] )
-          rf_dset.add( r.RooArgSet(rf_obs) )
+            for i in range(self.ndiscvars):
+                rf_obs.at(i).setVal(row[i])
+            rf_dset.add(r.RooArgSet(rf_obs))
 
         # pdfs should now be of the RooFit form
         rf_pdfs = r.RooArgList()
         for pdf in self.pdfs:
-            if not pdf.InheritsFrom('RooAbsPdf'):
-                raise RuntimeError('Found a pdf which does not inherit from RooAbsPdf.')
-            #pdf.getParameters( rf_dset ).setAttribAll("Constant")
-            rf_pdfs.add( pdf )
+            if not pdf.InheritsFrom("RooAbsPdf"):
+                raise RuntimeError("Found a pdf which does not inherit from RooAbsPdf.")
+            # pdf.getParameters( rf_dset ).setAttribAll("Constant")
+            rf_pdfs.add(pdf)
 
         # yields can still be numbers
-        rfylds = [ r.RooRealVar('y_%s'%self.compnames[i], 'y_%s'%self.compnames[i], yld, 0., 2.5*yld ) for i, yld in enumerate(self.yields) ]
+        rfylds = [
+            r.RooRealVar(
+                "y_%s" % self.compnames[i],
+                "y_%s" % self.compnames[i],
+                yld,
+                0.0,
+                2.5 * yld,
+            )
+            for i, yld in enumerate(self.yields)
+        ]
         rf_ylds = r.RooArgList()
         for yld in rfylds:
-          rf_ylds.add( yld  )
+            rf_ylds.add(yld)
 
         # now can create the pdf
-        rf_totpdf = r.RooAddPdf( 'pdf', 'pdf', rf_pdfs, rf_ylds, False )
+        rf_totpdf = r.RooAddPdf("pdf", "pdf", rf_pdfs, rf_ylds, False)
 
         # fit it (probably already been done)
-        rf_totpdf.fitTo( rf_dset, rf.Extended(True), rf.PrintLevel(-1) )
+        rf_totpdf.fitTo(rf_dset, r.RooFit.Extended(True), r.RooFit.PrintLevel(-1))
 
         # plot it
         c = r.TCanvas()
         for obs in self.rfobs:
-          pl = obs.frame()
-          rf_dset.plotOn(pl, rf.Binning(50) )
-          rf_totpdf.plotOn(pl)
-          pl.Draw()
-          c.Update()
-          c.Modified()
-          c.Draw()
-          r.gErrorIgnoreLevel = r.kInfo
-          c.Print('figs/rf_%s.pdf'%obs.GetName())
+            pl = obs.frame()
+            rf_dset.plotOn(pl, r.RooFit.Binning(50))
+            rf_totpdf.plotOn(pl)
+            pl.Draw()
+            c.Update()
+            c.Modified()
+            c.Draw()
+            r.gErrorIgnoreLevel = r.kInfo
+            c.Print("figs/rf_%s.pdf" % obs.GetName())
 
         # now get the sweights using RooStats
-        rf_sp = rs.SPlot( 'sdata', 'sdata', rf_dset, rf_totpdf, rf_ylds )
+        r.RooStats.SPlot("sdata", "sdata", rf_dset, rf_totpdf, rf_ylds)
 
-        weights = np.zeros( (rf_dset.numEntries(), self.ncomps) )
+        weights = np.zeros((rf_dset.numEntries(), self.ncomps))
 
         for ev in range(rf_dset.numEntries()):
-          rfvals = rf_dset.get(ev)
-          for i, obs in enumerate(self.rfobs):
-              assert( abs(rfvals.getRealValue( obs.GetName() ) - self.data[ev][i])<1e-6 )
-          for i, comp in enumerate(self.compnames):
-            weights[ev,i] = rfvals.getRealValue( 'y_'+comp+'_sw' )
+            rfvals = rf_dset.get(ev)
+            for i, obs in enumerate(self.rfobs):
+                assert abs(rfvals.getRealValue(obs.GetName()) - self.data[ev][i]) < 1e-6
+            for i, comp in enumerate(self.compnames):
+                weights[ev, i] = rfvals.getRealValue("y_" + comp + "_sw")
 
-        data_w_weights = np.append( self.data, weights, axis=1 )
-        sorted_data_w_weights = data_w_weights[np.argsort(data_w_weights[:,0])]
+        data_w_weights = np.append(self.data, weights, axis=1)
+        sorted_data_w_weights = data_w_weights[np.argsort(data_w_weights[:, 0])]
 
         return sorted_data_w_weights
 
-def convertRooAbsPdf(pdf,obs,npoints=400,forcenorm=False):
+
+def convertRooAbsPdf(pdf, obs, npoints=400, forcenorm=False):
     """
-    Helper function to convert a RooFit::RooAbsPdf object into a python callable that can be
-    used by either the `sweight` or `cow` classes
+    Convert RooAbsPdf into python callable.
+
+    Helper function to convert a RooFit::RooAbsPdf object into a python
+    callable that can be used by either the `sweight` or `cow` classes
 
     Parameters
     ----------
     pdf : RooAbsPdf
-        The pdf, must inherit from RooAbsPdf (e.g. RooGaussian, RooExponential, RooAddPdf etc.)
+        The pdf, must inherit from RooAbsPdf (e.g. RooGaussian, RooExponential,
+        RooAddPdf etc.)
     obs : RooRealVar
-        The observable, must inherit from RooRealVarLValue but will usually be a RooRealVar
+        The observable, must inherit from RooRealVarLValue but will usually be
+        a RooRealVar
     npoints : int, optional
         The number of points to use for the interpolation
     forcenorm : bool, optional
-        Force the return function to be normalised by performing a numerical integration of it
-        (the function should in most cases be normalised properly anyway so this shouldn't be
-        needed)
+        Force the return function to be normalised by performing a numerical
+        integration of it (the function should in most cases be normalised
+        properly anyway so this shouldn't be needed)
 
     Returns
     -------
     callable :
-        A callable function representing a normalised pdf which can then be passed to the `sweight`
-        or `cow` classes
+        A callable function representing a normalised pdf which can then be
+        passed to the `sweight` or `cow` classes
 
     """
     try:
         import ROOT as r
-        from ROOT import RooFit as rf
-    except:
-        raise RuntimeError('ROOT and RooFit must be installed to convert a RooAbsPdf')
+    except Exception:
+        raise RuntimeError("ROOT and RooFit must be installed to convert a RooAbsPdf")
 
-    if not hasattr(obs,'InheritsFrom'):
-        raise RuntimeError('Observable does not appear to be a ROOT like object - it should inherit from RooAbsReal. Type: ', type(obs))
+    if not hasattr(obs, "InheritsFrom"):
+        raise RuntimeError(
+            """Observable does not appear to be a ROOT like object - it should
+            inherit from RooAbsReal. Type: """,
+            type(obs),
+        )
 
-    if not obs.InheritsFrom('RooAbsRealLValue'):
-        raise RuntimeError('Observable does not appear to be of the right type - it should inherit from RooAbsRealLValue. Type: ', type(obs))
+    if not obs.InheritsFrom("RooAbsRealLValue"):
+        raise RuntimeError(
+            """Observable does not appear to be of the right type - it should
+            inherit from RooAbsRealLValue. Type: """,
+            type(obs),
+        )
 
-    range = ( obs.getMin(), obs.getMax() )
+    range = (obs.getMin(), obs.getMax())
 
-    xvals = np.linspace(*range,npoints)
+    xvals = np.linspace(*range, npoints)
 
     yvals = []
 
     normset = r.RooArgSet(obs)
     for x in xvals:
         obs.setVal(x)
-        yvals.append( pdf.getVal(normset) )
+        yvals.append(pdf.getVal(normset))
 
-    f = InterpolatedUnivariateSpline(xvals,yvals)
+    f = InterpolatedUnivariateSpline(xvals, yvals)
 
     N = 1
     if forcenorm:
-        N = nquad(f,(range,))[0]
+        N = nquad(f, (range,))[0]
 
-    retf = lambda x: f(x)/N
+    def retf(x):
+        return f(x) / N
 
     return retf


Test patch:
diff --git a/tests/check_roofit_conversion.py b/tests/check_roofit_conversion.py
index 05799ff..6c86790 100644
--- a/tests/check_roofit_conversion.py
+++ b/tests/check_roofit_conversion.py
@@ -1,28 +1,41 @@
 import ROOT as r
 from ROOT import RooFit as rf
 
-mass = r.RooRealVar('m','m',5000,5800)
+mass = r.RooRealVar("m", "m", 5000, 5800)
 
 # gaussian
-mean = r.RooRealVar('mu','mu',5350,5300,5400)
-sigma = r.RooRealVar('sg','sg',25,0,100)
-spdf = r.RooGaussian('sig','sig',mass,mean,sigma)
+mean = r.RooRealVar("mu", "mu", 5350, 5300, 5400)
+sigma = r.RooRealVar("sg", "sg", 25, 0, 100)
+spdf = r.RooGaussian("sig", "sig", mass, mean, sigma)
 
 # exponential
-slope = r.RooRealVar('lb','lb',-0.002,-0.01,0)
-bpdf = r.RooExponential('bkg','bkg',mass,slope)
+slope = r.RooRealVar("lb", "lb", -0.002, -0.01, 0)
+bpdf = r.RooExponential("bkg", "bkg", mass, slope)
 
 # total
-sy = r.RooRealVar('sy','sy',1000,0,10000)
-by = r.RooRealVar('by','by',2000,0,10000)
-pdf = r.RooAddPdf('pdf','pdf',r.RooArgList(spdf,bpdf), r.RooArgList(sy,by))
+sy = r.RooRealVar("sy", "sy", 1000, 0, 10000)
+by = r.RooRealVar("by", "by", 2000, 0, 10000)
+pdf = r.RooAddPdf("pdf", "pdf", r.RooArgList(spdf, bpdf), r.RooArgList(sy, by))
 
 # plot RooFit style
 c = r.TCanvas()
 pl = mass.frame(rf.Bins(800))
-pdf.plotOn(pl, rf.Normalization(1,r.RooAbsReal.NumEvent), rf.Precision(-1), rf.Components("bkg"), rf.LineColor(r.kRed), rf.LineStyle(r.kDashed) )
-pdf.plotOn(pl, rf.Normalization(1,r.RooAbsReal.NumEvent), rf.Precision(-1), rf.Components("sig"), rf.LineColor(r.kGreen+3) )
-pdf.plotOn(pl, rf.Normalization(1,r.RooAbsReal.NumEvent), rf.Precision(-1) )
+pdf.plotOn(
+    pl,
+    rf.Normalization(1, r.RooAbsReal.NumEvent),
+    rf.Precision(-1),
+    rf.Components("bkg"),
+    rf.LineColor(r.kRed),
+    rf.LineStyle(r.kDashed),
+)
+pdf.plotOn(
+    pl,
+    rf.Normalization(1, r.RooAbsReal.NumEvent),
+    rf.Precision(-1),
+    rf.Components("sig"),
+    rf.LineColor(r.kGreen + 3),
+)
+pdf.plotOn(pl, rf.Normalization(1, r.RooAbsReal.NumEvent), rf.Precision(-1))
 pl.Draw()
 c.Update()
 
@@ -30,21 +43,20 @@
 # now do the conversion
 # this should return every pdf normalised
 from sweights import convertRooAbsPdf
-_spdf = convertRooAbsPdf(spdf,mass)
-_bpdf = convertRooAbsPdf(bpdf,mass)
-_pdf = convertRooAbsPdf(pdf,mass)
+
+_spdf = convertRooAbsPdf(spdf, mass)
+_bpdf = convertRooAbsPdf(bpdf, mass)
+_pdf = convertRooAbsPdf(pdf, mass)
 
 # now draw the function mpl style
 import matplotlib.pyplot as plt
 import numpy as np
-fs = sy.getVal()/(sy.getVal()+by.getVal())
-fb = by.getVal()/(sy.getVal()+by.getVal())
+
+fs = sy.getVal() / (sy.getVal() + by.getVal())
+fb = by.getVal() / (sy.getVal() + by.getVal())
 
 x = np.linspace(mass.getMin(), mass.getMax(), 400)
-plt.plot(x, fs*_spdf(x), 'g-')
-plt.plot(x, fb*_bpdf(x), 'r--')
-plt.plot(x, _pdf(x), 'b-')
+plt.plot(x, fs * _spdf(x), "g-")
+plt.plot(x, fb * _bpdf(x), "r--")
+plt.plot(x, _pdf(x), "b-")
 plt.show()
-
-
-
diff --git a/tests/example.py b/tests/example.py
index eefc947..527518e 100644
--- a/tests/example.py
+++ b/tests/example.py
@@ -19,289 +19,372 @@
 
 Ns = 5000
 Nb = 5000
-ypars = [Ns,Nb]
+ypars = [Ns, Nb]
 
 # mass
-mrange = (0,1)
+mrange = (0, 1)
 mu = 0.5
 sg = 0.1
 lb = 10
-mpars = [mu,sg,lb]
+mpars = [mu, sg, lb]
 
 # decay time
-trange = (0,1)
+trange = (0, 1)
 tlb = 2
 tpars = [tlb]
 
 # generate the toy
-def generate(Ns,Nb,mu,sg,lb,tlb,poisson=False,ret_true=False):
+def generate(Ns, Nb, mu, sg, lb, tlb, poisson=False, ret_true=False):
 
-  Nsig = np.random.poisson(Ns) if poisson else Ns
-  Nbkg = np.random.poisson(Nb) if poisson else Nb
+    Nsig = np.random.poisson(Ns) if poisson else Ns
+    Nbkg = np.random.poisson(Nb) if poisson else Nb
 
-  sigM = norm(mu,sg)
-  bkgM = expon(mrange[0], lb)
+    sigM = norm(mu, sg)
+    bkgM = expon(mrange[0], lb)
 
-  sigT = expon(trange[0], tlb)
-  bkgT = uniform(trange[0],trange[1]-trange[0])
+    sigT = expon(trange[0], tlb)
+    bkgT = uniform(trange[0], trange[1] - trange[0])
 
-  # generate
-  sigMflt = sigM.cdf(mrange)
-  bkgMflt = bkgM.cdf(mrange)
-  sigTflt = sigT.cdf(trange)
-  bkgTflt = bkgT.cdf(trange)
+    # generate
+    sigMflt = sigM.cdf(mrange)
+    bkgMflt = bkgM.cdf(mrange)
+    sigTflt = sigT.cdf(trange)
+    bkgTflt = bkgT.cdf(trange)
 
-  sigMvals = sigM.ppf( np.random.uniform(*sigMflt,size=Nsig) )
-  sigTvals = sigT.ppf( np.random.uniform(*sigTflt,size=Nsig) )
+    sigMvals = sigM.ppf(np.random.uniform(*sigMflt, size=Nsig))
+    sigTvals = sigT.ppf(np.random.uniform(*sigTflt, size=Nsig))
 
-  bkgMvals = bkgM.ppf( np.random.uniform(*bkgMflt,size=Nbkg) )
-  bkgTvals = bkgT.ppf( np.random.uniform(*bkgTflt,size=Nbkg) )
+    bkgMvals = bkgM.ppf(np.random.uniform(*bkgMflt, size=Nbkg))
+    bkgTvals = bkgT.ppf(np.random.uniform(*bkgTflt, size=Nbkg))
 
-  Mvals = np.concatenate( (sigMvals, bkgMvals) )
-  Tvals = np.concatenate( (sigTvals, bkgTvals) )
+    Mvals = np.concatenate((sigMvals, bkgMvals))
+    Tvals = np.concatenate((sigTvals, bkgTvals))
 
-  truth = np.concatenate( ( np.ones_like(sigMvals), np.zeros_like(bkgMvals) ) )
+    truth = np.concatenate((np.ones_like(sigMvals), np.zeros_like(bkgMvals)))
+
+    if ret_true:
+        return np.stack((Mvals, Tvals, truth), axis=1)
+    else:
+        return np.stack((Mvals, Tvals), axis=1)
 
-  if ret_true:
-    return np.stack( (Mvals,Tvals,truth), axis=1 )
-  else:
-    return np.stack( (Mvals,Tvals), axis=1 )
 
 # define mass pdf for plotting etc.
-def mpdf(x, Ns, Nb, mu, sg, lb, comps=['sig','bkg']):
+def mpdf(x, Ns, Nb, mu, sg, lb, comps=["sig", "bkg"]):
+
+    sig = norm(mu, sg)
+    sigN = np.diff(sig.cdf(mrange))
 
-  sig  = norm(mu,sg)
-  sigN = np.diff( sig.cdf(mrange) )
+    bkg = expon(mrange[0], lb)
+    bkgN = np.diff(bkg.cdf(mrange))
 
-  bkg  = expon(mrange[0], lb)
-  bkgN = np.diff( bkg.cdf(mrange) )
+    tot = 0
+    if "sig" in comps:
+        tot += Ns * sig.pdf(x) / sigN
+    if "bkg" in comps:
+        tot += Nb * bkg.pdf(x) / bkgN
 
-  tot = 0
-  if 'sig' in comps: tot += Ns * sig.pdf(x) / sigN
-  if 'bkg' in comps: tot += Nb * bkg.pdf(x) / bkgN
+    return tot
 
-  return tot
 
 # define mass pdf for iminuit fitting
 def mpdf_min(x, Ns, Nb, mu, sg, lb):
-  return (Ns+Nb, mpdf(x, Ns, Nb, mu, sg, lb) )
+    return (Ns + Nb, mpdf(x, Ns, Nb, mu, sg, lb))
+
 
 # define time pdf for plotting etc.
-def tpdf(x, Ns, Nb, tlb, comps=['sig','bkg']):
+def tpdf(x, Ns, Nb, tlb, comps=["sig", "bkg"]):
 
-  sig  = expon(trange[0],tlb)
-  sigN = np.diff( sig.cdf(trange) )
+    sig = expon(trange[0], tlb)
+    sigN = np.diff(sig.cdf(trange))
 
-  bkg  = uniform(trange[0],trange[1]-trange[0])
-  bkgN = np.diff( bkg.cdf(trange) )
+    bkg = uniform(trange[0], trange[1] - trange[0])
+    bkgN = np.diff(bkg.cdf(trange))
 
-  tot = 0
-  if 'sig' in comps: tot += Ns * sig.pdf(x) / sigN
-  if 'bkg' in comps: tot += Nb * bkg.pdf(x) / bkgN
+    tot = 0
+    if "sig" in comps:
+        tot += Ns * sig.pdf(x) / sigN
+    if "bkg" in comps:
+        tot += Nb * bkg.pdf(x) / bkgN
+
+    return tot
 
-  return tot
 
 # define signal only time pdf for cov corrector
 def tpdf_cor(x, tlb):
-  return tpdf(x,1,0,tlb,['sig'])
+    return tpdf(x, 1, 0, tlb, ["sig"])
+
 
 def myerrorbar(data, ax, bins, range, wts=None, label=None, col=None):
-  col = col or 'k'
-  nh, xe = np.histogram(data,bins=bins,range=range)
-  cx = 0.5*(xe[1:]+xe[:-1])
-  err = nh**0.5
-  if wts is not None:
-    whist = bh.Histogram( bh.axis.Regular(bins,*range), storage=bh.storage.Weight() )
-    whist.fill( data, weight = wts )
-    cx = whist.axes[0].centers
-    nh = whist.view().value
-    err = whist.view().variance**0.5
-
-  ax.errorbar(cx, nh, err, capsize=2,label=label,fmt=f'{col}o')
+    col = col or "k"
+    nh, xe = np.histogram(data, bins=bins, range=range)
+    cx = 0.5 * (xe[1:] + xe[:-1])
+    err = nh ** 0.5
+    if wts is not None:
+        whist = bh.Histogram(bh.axis.Regular(bins, *range), storage=bh.storage.Weight())
+        whist.fill(data, weight=wts)
+        cx = whist.axes[0].centers
+        nh = whist.view().value
+        err = whist.view().variance ** 0.5
+
+    ax.errorbar(cx, nh, err, capsize=2, label=label, fmt=f"{col}o")
+
 
 def plot(toy, draw_pdf=True, save=None):
 
-  nbins = 50
+    nbins = 50
+
+    fig, ax = plt.subplots(1, 2, figsize=(12, 4))
 
-  fig, ax = plt.subplots(1,2,figsize=(12,4))
+    myerrorbar(toy[:, 0], ax[0], bins=nbins, range=mrange)
+    myerrorbar(toy[:, 1], ax[1], bins=nbins, range=trange)
 
-  myerrorbar(toy[:,0],ax[0],bins=nbins,range=mrange)
-  myerrorbar(toy[:,1],ax[1],bins=nbins,range=trange)
+    if draw_pdf:
+        m = np.linspace(*mrange, 400)
+        mN = (mrange[1] - mrange[0]) / nbins
 
-  if draw_pdf:
-    m = np.linspace(*mrange,400)
-    mN = (mrange[1]-mrange[0])/nbins
+        bkgm = mpdf(m, *(ypars + mpars), comps=["bkg"])
+        sigm = mpdf(m, *(ypars + mpars), comps=["sig"])
+        totm = bkgm + sigm
 
-    bkgm = mpdf(m, *(ypars+mpars),comps=['bkg'])
-    sigm = mpdf(m, *(ypars+mpars),comps=['sig'])
-    totm = bkgm + sigm
+        ax[0].plot(m, mN * bkgm, "r--", label="Background")
+        ax[0].plot(m, mN * sigm, "g:", label="Signal")
+        ax[0].plot(m, mN * totm, "b-", label="Total PDF")
 
-    ax[0].plot(m, mN*bkgm, 'r--', label='Background')
-    ax[0].plot(m, mN*sigm, 'g:' , label='Signal')
-    ax[0].plot(m, mN*totm, 'b-' , label='Total PDF')
+        t = np.linspace(*trange, 400)
+        tN = (trange[1] - trange[0]) / nbins
 
-    t = np.linspace(*trange,400)
-    tN = (trange[1]-trange[0])/nbins
+        bkgt = tpdf(t, *(ypars + tpars), comps=["bkg"])
+        sigt = tpdf(t, *(ypars + tpars), comps=["sig"])
+        tott = bkgt + sigt
 
-    bkgt = tpdf(t, *(ypars+tpars),comps=['bkg'])
-    sigt = tpdf(t, *(ypars+tpars),comps=['sig'])
-    tott = bkgt + sigt
+        ax[1].plot(t, tN * bkgt, "r--", label="Background")
+        ax[1].plot(t, tN * sigt, "g:", label="Signal")
+        ax[1].plot(t, tN * tott, "b-", label="Total PDF")
 
-    ax[1].plot(t, tN*bkgt, 'r--', label='Background')
-    ax[1].plot(t, tN*sigt, 'g:' , label='Signal')
-    ax[1].plot(t, tN*tott, 'b-' , label='Total PDF')
+    ax[0].set_xlabel("Mass")
+    ax[0].set_ylim(bottom=0)
+    ax[0].legend()
 
-  ax[0].set_xlabel('Mass')
-  ax[0].set_ylim(bottom=0)
-  ax[0].legend()
+    ax[1].set_xlabel("Time")
+    ax[1].set_ylim(bottom=0)
+    ax[1].legend()
 
-  ax[1].set_xlabel('Time')
-  ax[1].set_ylim(bottom=0)
-  ax[1].legend()
+    fig.tight_layout()
+    if save:
+        fig.savefig(save)
 
-  fig.tight_layout()
-  if save: fig.savefig(save)
 
-def plot_wts(x, sw, bw, ylabel='Weight',save=None):
+def plot_wts(x, sw, bw, ylabel="Weight", save=None):
+
+    fig, ax = plt.subplots()
+    ax.plot(x, sw, "b--", label="Signal")
+    ax.plot(x, bw, "r:", label="Background")
+    ax.plot(x, sw + bw, "k-", label="Sum")
+    ax.set_xlabel("Mass")
+    ax.set_ylabel("Weight")
+    fig.tight_layout()
+    if save:
+        fig.savefig(save)
 
-  fig,ax = plt.subplots()
-  ax.plot(x, sw, 'b--', label='Signal')
-  ax.plot(x, bw, 'r:' , label='Background')
-  ax.plot(x, sw+bw, 'k-', label='Sum')
-  ax.set_xlabel('Mass')
-  ax.set_ylabel('Weight')
-  fig.tight_layout()
-  if save: fig.savefig(save)
 
 def plot_tweighted(x, wts, wtnames=[], funcs=[], save=None):
 
-  fig, ax = plt.subplots()
+    fig, ax = plt.subplots()
+
+    t = np.linspace(*trange, 400)
+    N = (trange[1] - trange[0]) / 50
 
-  t = np.linspace(*trange,400)
-  N = (trange[1]-trange[0])/50
+    for i, wt in enumerate(wts):
+        label = None
+        if i < len(wtnames):
+            label = wtnames[i]
+        myerrorbar(x, ax, bins=50, range=trange, wts=wt, label=label, col=f"C{i}")
+        if i < len(funcs):
+            ax.plot(t, N * funcs[i](t), f"C{i}-")
 
-  for i, wt in enumerate(wts):
-    label = None
-    if i<len(wtnames): label = wtnames[i]
-    myerrorbar(x, ax, bins=50, range=trange, wts=wt, label=label, col=f'C{i}')
-    if i<len(funcs):
-      ax.plot(t,N*funcs[i](t),f'C{i}-')
+    ax.legend()
+    ax.set_xlabel("Time")
+    ax.set_ylabel("Weighted Events")
+    fig.tight_layout()
+    if save:
+        fig.savefig(save)
 
-  ax.legend()
-  ax.set_xlabel('Time')
-  ax.set_ylabel('Weighted Events')
-  fig.tight_layout()
-  if save: fig.savefig(save)
 
 def wnll(tlb, tdata, wts):
-  sig  = expon(trange[0],tlb)
-  sigN = np.diff( sig.cdf(trange) )
-  return -np.sum( wts * np.log( sig.pdf( tdata ) / sigN ) )
-
-if __name__ == '__main__':
-
-  parser = ArgumentParser()
-  parser.add_argument('-p','--makeplots'  , default=False, action='store_true', help='Make plots')
-  parser.add_argument('-d','--plotdir'    , default='__plots__'               , help='Save location for plots')
-  parser.add_argument('-s','--seed'       , default=None, type=int            , help='Set random seed')
-  parser.add_argument('-i','--interactive', default=False, action='store_true', help='Show plots interactively at the end')
-  parser.add_argument('-v','--verbose'    , default=False, action='store_true', help='Print more output')
-  args = parser.parse_args()
-
-  if args.seed:
-    np.random.seed(args.seed)
-
-  if args.makeplots and not os.path.exists(args.plotdir):
-    os.system(f'mkdir -p {args.plotdir}')
-
-  ## generate the toy
-  toy = generate(Ns,Nb,mu,sg,lb,tlb,ret_true=True)
-  if args.makeplots: plot(toy, save=f'{args.plotdir}/toy.png')
-
-  ## print kendall rank coeff
-  kts = kendall_tau(toy[:,0],toy[:,1])
-  print('Kendall Tau:', pdg_format( kts[0], kts[1] ) )
-
-  ## fit the toy mass
-  mi = Minuit( ExtendedUnbinnedNLL(toy[:,0], mpdf_min), Ns=Ns, Nb=Nb, mu=mu, sg=sg, lb=lb )
-  mi.limits['Ns'] = (0,Ns+Nb)
-  mi.limits['Nb'] = (0,Ns+Nb)
-  mi.limits['mu'] = mrange
-  mi.limits['sg'] = (0,mrange[1]-mrange[0])
-  mi.limits['lb'] = (0,50)
-
-  mi.migrad()
-  mi.hesse()
-  if args.verbose: print(mi)
-
-  # define estimated functions
-  spdf = lambda m: mpdf(m,*mi.values,comps=['sig'])
-  bpdf = lambda m: mpdf(m,*mi.values,comps=['bkg'])
-
-  # make the sweighter
-  print('Compute sweights')
-  sweighter = sweight( toy[:,0], [spdf,bpdf], [mi.values['Ns'],mi.values['Nb']], (mrange,), method='summation', compnames=('sig','bkg'), verbose=args.verbose, checks=args.verbose )
-
-  # get the COW equivalents (they should be normalised but the cow should also take care if they are not)
-  gs = lambda m: mpdf(m,*mi.values,comps=['sig']) / mi.values['Ns']
-  gb = lambda m: mpdf(m,*mi.values,comps=['bkg']) / mi.values['Nb']
-  Im = 1 #lambda m: mpdf(m,*mi.values) / (mi.values['Ns'] + mi.values['Nb'] )
-
-  # make the cow
-  print('Compute cow weights')
-  cw = cow(mrange, spdf, gb, Im, verbose=args.verbose)
-
-  # compare the two
-  flbs = []
-  for meth, cls in zip( ['SW','COW'], [sweighter,cw] ):
-
-    # plot weights
-    x = np.linspace(*mrange,400)
-    swp = cls.getWeight(0,x)
-    bwp = cls.getWeight(1,x)
-    if args.makeplots: plot_wts(x, swp, bwp, save=f'{args.plotdir}/{meth}_wts.png')
-
-    # fit weighted data
-    wts = cls.getWeight(0,toy[:,0])
-    nll = lambda tlb: wnll(tlb, toy[:,1], wts)
-
-    # do the minimisation
-    tmi = Minuit( nll, tlb=tlb )
-    tmi.limits['tlb'] = (1,3)
-    tmi.errordef = Minuit.LIKELIHOOD
-    tmi.migrad()
-    tmi.hesse()
-
-    # and do the correction
-    fval = np.array(tmi.values)
-    flbs.append(fval[0])
-    fcov = np.array( tmi.covariance.tolist() )
-
-    # first order correction
-    ncov = approx_cov_correct(tpdf_cor, toy[:,1], wts, fval, fcov, verbose=args.verbose)
-
-    # second order correction
-    hs  = tpdf_cor
-    ws  = lambda m: cls.getWeight(0,m)
-    W   = cls.Wkl
-
-    # these derivatives can be done numerically but for the sweights / COW case it's straightfoward to compute them
-    ws = lambda Wss, Wsb, Wbb, gs, gb: (Wbb*gs - Wsb*gb) / ((Wbb-Wsb)*gs + (Wss-Wsb)*gb)
-    dws_Wss = lambda Wss, Wsb, Wbb, gs, gb: gb * ( Wsb*gb - Wbb*gs ) / (-Wss*gb + Wsb*gs + Wsb*gb - Wbb*gs)**2
-    dws_Wsb = lambda Wss, Wsb, Wbb, gs, gb: ( Wbb*gs**2 - Wss*gb**2 ) / (Wss*gb - Wsb*gs - Wsb*gb + Wbb*gs)**2
-    dws_Wbb = lambda Wss, Wsb, Wbb, gs, gb: gs * ( Wss*gb - Wsb*gs ) / (-Wss*gb + Wsb*gs + Wsb*gb - Wbb*gs)**2
-
-    tcov = cov_correct(hs, [gs,gb], toy[:,1], toy[:,0], wts, [mi.values['Ns'],mi.values['Nb']], fval, fcov, [dws_Wss,dws_Wsb,dws_Wbb],[W[0,0],W[0,1],W[1,1]], verbose=args.verbose)
-
-    if args.verbose: print('Method:', meth, f'- covariance corrected {fval[0]:.1f} +/- {fcov[0,0]**0.5:.1f} ---> {fval[0]:.1f} +/- {tcov[0,0]**0.5:.1f}')
-
-  ## plot weight T distribution
-  swf  = lambda t: tpdf(t, mi.values['Ns'], 0, flbs[0], comps=['sig'] )
-  cowf = lambda t: tpdf(t, mi.values['Ns'], 0, flbs[1], comps=['sig'] )
-  sws  = sweighter.getWeight(0, toy[:,0])
-  scow = cw.getWeight(0, toy[:,0])
-
-  if args.makeplots: plot_tweighted(toy[:,1], [sws,scow], ['SW','COW'], funcs=[swf,cowf], save=f'{args.plotdir}/tfit.png' )
-
-
-  if args.interactive: plt.show()
+    sig = expon(trange[0], tlb)
+    sigN = np.diff(sig.cdf(trange))
+    return -np.sum(wts * np.log(sig.pdf(tdata) / sigN))
+
+
+if __name__ == "__main__":
+
+    parser = ArgumentParser()
+    parser.add_argument(
+        "-p", "--makeplots", default=False, action="store_true", help="Make plots"
+    )
+    parser.add_argument(
+        "-d", "--plotdir", default="__plots__", help="Save location for plots"
+    )
+    parser.add_argument("-s", "--seed", default=None, type=int, help="Set random seed")
+    parser.add_argument(
+        "-i",
+        "--interactive",
+        default=False,
+        action="store_true",
+        help="Show plots interactively at the end",
+    )
+    parser.add_argument(
+        "-v", "--verbose", default=False, action="store_true", help="Print more output"
+    )
+    args = parser.parse_args()
+
+    if args.seed:
+        np.random.seed(args.seed)
+
+    if args.makeplots and not os.path.exists(args.plotdir):
+        os.system(f"mkdir -p {args.plotdir}")
+
+    ## generate the toy
+    toy = generate(Ns, Nb, mu, sg, lb, tlb, ret_true=True)
+    if args.makeplots:
+        plot(toy, save=f"{args.plotdir}/toy.png")
+
+    ## print kendall rank coeff
+    kts = kendall_tau(toy[:, 0], toy[:, 1])
+    print("Kendall Tau:", pdg_format(kts[0], kts[1]))
+
+    ## fit the toy mass
+    mi = Minuit(
+        ExtendedUnbinnedNLL(toy[:, 0], mpdf_min), Ns=Ns, Nb=Nb, mu=mu, sg=sg, lb=lb
+    )
+    mi.limits["Ns"] = (0, Ns + Nb)
+    mi.limits["Nb"] = (0, Ns + Nb)
+    mi.limits["mu"] = mrange
+    mi.limits["sg"] = (0, mrange[1] - mrange[0])
+    mi.limits["lb"] = (0, 50)
+
+    mi.migrad()
+    mi.hesse()
+    if args.verbose:
+        print(mi)
+
+    # define estimated functions
+    spdf = lambda m: mpdf(m, *mi.values, comps=["sig"])
+    bpdf = lambda m: mpdf(m, *mi.values, comps=["bkg"])
+
+    # make the sweighter
+    print("Compute sweights")
+    sweighter = sweight(
+        toy[:, 0],
+        [spdf, bpdf],
+        [mi.values["Ns"], mi.values["Nb"]],
+        (mrange,),
+        method="summation",
+        compnames=("sig", "bkg"),
+        verbose=args.verbose,
+        checks=args.verbose,
+    )
+
+    # get the COW equivalents (they should be normalised but the cow should also take care if they are not)
+    gs = lambda m: mpdf(m, *mi.values, comps=["sig"]) / mi.values["Ns"]
+    gb = lambda m: mpdf(m, *mi.values, comps=["bkg"]) / mi.values["Nb"]
+    Im = 1  # lambda m: mpdf(m,*mi.values) / (mi.values['Ns'] + mi.values['Nb'] )
+
+    # make the cow
+    print("Compute cow weights")
+    cw = cow(mrange, spdf, gb, Im, verbose=args.verbose)
+
+    # compare the two
+    flbs = []
+    for meth, cls in zip(["SW", "COW"], [sweighter, cw]):
+
+        # plot weights
+        x = np.linspace(*mrange, 400)
+        swp = cls.getWeight(0, x)
+        bwp = cls.getWeight(1, x)
+        if args.makeplots:
+            plot_wts(x, swp, bwp, save=f"{args.plotdir}/{meth}_wts.png")
+
+        # fit weighted data
+        wts = cls.getWeight(0, toy[:, 0])
+        nll = lambda tlb: wnll(tlb, toy[:, 1], wts)
+
+        # do the minimisation
+        tmi = Minuit(nll, tlb=tlb)
+        tmi.limits["tlb"] = (1, 3)
+        tmi.errordef = Minuit.LIKELIHOOD
+        tmi.migrad()
+        tmi.hesse()
+
+        # and do the correction
+        fval = np.array(tmi.values)
+        flbs.append(fval[0])
+        fcov = np.array(tmi.covariance.tolist())
+
+        # first order correction
+        ncov = approx_cov_correct(
+            tpdf_cor, toy[:, 1], wts, fval, fcov, verbose=args.verbose
+        )
+
+        # second order correction
+        hs = tpdf_cor
+        ws = lambda m: cls.getWeight(0, m)
+        W = cls.Wkl
+
+        # these derivatives can be done numerically but for the sweights / COW case it's straightfoward to compute them
+        ws = lambda Wss, Wsb, Wbb, gs, gb: (Wbb * gs - Wsb * gb) / (
+            (Wbb - Wsb) * gs + (Wss - Wsb) * gb
+        )
+        dws_Wss = (
+            lambda Wss, Wsb, Wbb, gs, gb: gb
+            * (Wsb * gb - Wbb * gs)
+            / (-Wss * gb + Wsb * gs + Wsb * gb - Wbb * gs) ** 2
+        )
+        dws_Wsb = (
+            lambda Wss, Wsb, Wbb, gs, gb: (Wbb * gs ** 2 - Wss * gb ** 2)
+            / (Wss * gb - Wsb * gs - Wsb * gb + Wbb * gs) ** 2
+        )
+        dws_Wbb = (
+            lambda Wss, Wsb, Wbb, gs, gb: gs
+            * (Wss * gb - Wsb * gs)
+            / (-Wss * gb + Wsb * gs + Wsb * gb - Wbb * gs) ** 2
+        )
+
+        tcov = cov_correct(
+            hs,
+            [gs, gb],
+            toy[:, 1],
+            toy[:, 0],
+            wts,
+            [mi.values["Ns"], mi.values["Nb"]],
+            fval,
+            fcov,
+            [dws_Wss, dws_Wsb, dws_Wbb],
+            [W[0, 0], W[0, 1], W[1, 1]],
+            verbose=args.verbose,
+        )
+
+        if args.verbose:
+            print(
+                "Method:",
+                meth,
+                f"- covariance corrected {fval[0]:.1f} +/- {fcov[0,0]**0.5:.1f} ---> {fval[0]:.1f} +/- {tcov[0,0]**0.5:.1f}",
+            )
+
+    ## plot weight T distribution
+    swf = lambda t: tpdf(t, mi.values["Ns"], 0, flbs[0], comps=["sig"])
+    cowf = lambda t: tpdf(t, mi.values["Ns"], 0, flbs[1], comps=["sig"])
+    sws = sweighter.getWeight(0, toy[:, 0])
+    scow = cw.getWeight(0, toy[:, 0])
+
+    if args.makeplots:
+        plot_tweighted(
+            toy[:, 1],
+            [sws, scow],
+            ["SW", "COW"],
+            funcs=[swf, cowf],
+            save=f"{args.plotdir}/tfit.png",
+        )
+
+    if args.interactive:
+        plt.show()