update Bands for supporting n bands def

Author: TomDonoghue

specparam/bands/bands.py

@@ -20,19 +20,34 @@ class Bands():
     >>> bands = Bands({'theta' : [4, 8], 'alpha' : [8, 12], 'beta' : [15, 30]})
     """
 
-    def __init__(self, input_bands={}):
+    def __init__(self, input_bands=None, n_bands=None):
         """Initialize the Bands object.
 
         Parameters
         ----------
         input_bands : dict, optional
             A dictionary of oscillation bands.
+        n_bands : int, optional
+            The number of bands to extract from the spectra.
+            Can only be specified if not providing `input_bands`.
+
+        Attributes
+        ----------
+        bands : OrderedDict
+            Band definitions.
         """
 
         self.bands = OrderedDict()
 
-        for label, band_def in input_bands.items():
-            self.add_band(label, band_def)
+        if input_bands:
+            for label, band_def in input_bands.items():
+                self.add_band(label, band_def)
+
+        self._n_bands = None
+        if n_bands:
+            if input_bands:
+                raise ValueError('Cannot provive both `input_bands` and `n_bands`.')
+            self._n_bands = n_bands
 
 
     def __getitem__(self, label):
@@ -89,7 +104,10 @@ def definitions(self):
     def n_bands(self):
         """The number of bands defined in the object."""
 
-        return len(self.bands)
+        if self._n_bands is not None:
+            return self._n_bands
+        else:
+            return len(self.bands)
 
 
     def add_band(self, label, band_definition):
@@ -103,6 +121,7 @@ def add_band(self, label, band_definition):
             The lower and upper frequency limit of the band, in Hz.
         """
 
+        self._n_bands = None
         self._check_band(label, band_definition)
         self.bands[label] = tuple(band_definition)
 
@@ -147,3 +166,28 @@ def _check_band(label, band_definition):
         # Safety check that limits are in correct order
         if not band_definition[0] < band_definition[1]:
             raise ValueError("Band limit definitions are invalid.")
+
+
+def check_bands(bands):
+    """Check bands definition.
+
+    Parameters
+    ----------
+    bands : Bands or dict or int, optional
+        How to organize peaks into bands.
+
+    Returns
+    -------
+    bands : Bands
+        Bands definition.
+    """
+
+    if not isinstance(bands, Bands):
+        if isinstance(bands, (dict, OrderedDict)):
+            bands = Bands(bands)
+        elif isinstance(bands, int):
+            bands = Bands(n_bands=bands)
+        else:
+            raise ValueError('Bands definition not understood.')
+
+    return bands

specparam/tests/bands/test_bands.py

@@ -52,3 +52,33 @@ def test_bands_properties(tbands):
 
     assert set(tbands.labels) == set(['theta', 'alpha', 'beta'])
     assert tbands.n_bands == 3
+
+def test_bands_n_bands():
+
+    n_bands = 2
+    bands = Bands(n_bands=n_bands)
+    assert bands.bands == {}
+    assert bands._n_bands == n_bands
+    assert len(bands) == n_bands
+
+    # Check that adding a band replaces n_band definition
+    bands.add_band('alpha', [7, 14])
+    assert bands._n_bands == None
+    assert len(bands) == 1
+
+    # test fails adding both bands and n_bands
+    with raises(ValueError):
+        bands = Bands({'alpha' : (7, 14)}, n_bands=2)
+
+
+def test_check_bands(tbands):
+
+    out1 = check_bands(tbands)
+    assert isinstance(out1, Bands)
+    assert out1 == tbands
+
+    out2 = check_bands({'alpha' : (7, 14)})
+    assert isinstance(out2, Bands)
+
+    out3 = check_bands(2)
+    assert isinstance(out3, Bands)