Repository: mir-dataset-loaders/mirdata

Issue:
Multitrack Datasets
How can we best support multitrack datasets? The current solution is to have each 'Track' be a multitrack with a ton of attributes, but it's clunky and difficult to take full advantage of the multitracks themselves. We're also loosely tying a `Track` to something that can be mapped to a `jams` file/object, and jams isnt built for multiple audio files.

My current thinking is we could index at the stem level, and support a new base class `MultiTrack` which would group multiple `Track` objects. The grouping could be stored as part of the index, and any mixture-level annotations could itself be stored as a `Track` object, and be a part of the multitrack object.

Thoughts?

cc @nkundiushuti  - we could test this out with Phenicx-Anechoic #270 

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/mir-dataset-loaders/mirdata/pull/295?src=pr&el=h1) Report
> Merging [#295](https://codecov.io/gh/mir-dataset-loaders/mirdata/pull/295?src=pr&el=desc) into [master](https://codecov.io/gh/mir-dataset-loaders/mirdata/commit/1eedb64932c593b213e8c2218c1f35a0ee59f0b1?el=desc) will **decrease** coverage by `0.04%`.
> The diff coverage is `100.00%`.


```diff
@@            Coverage Diff             @@
##           master     #295      +/-   ##
==========================================
- Coverage   99.03%   98.98%   -0.05%     
==========================================
  Files          22       21       -1     
  Lines        2373     2263     -110     
==========================================
- Hits         2350     2240     -110     
  Misses         23       23              
```

nkundiushuti: > @nkundiushuti should be ready to go, let me know if this will work for #270 or if there's something missing

thanks! I will check this! 
rabitt: Merging!

Patch:
diff --git a/mirdata/track.py b/mirdata/track.py
index a503ea7..208174d 100644
--- a/mirdata/track.py
+++ b/mirdata/track.py
@@ -1,18 +1,18 @@
 # -*- coding: utf-8 -*-
 """track object utility functions
 """
-
-
 import types
 
+import numpy as np
+
 MAX_STR_LEN = 100
 
 
 class Track(object):
     def __repr__(self):
-        properties = [v for v in dir(self.__class__) if not v.startswith('_')]
+        properties = [v for v in dir(self.__class__) if not v.startswith("_")]
         attributes = [
-            v for v in dir(self) if not v.startswith('_') and v not in properties
+            v for v in dir(self) if not v.startswith("_") and v not in properties
         ]
 
         repr_str = "Track(\n"
@@ -21,7 +21,7 @@ class Track(object):
             val = getattr(self, attr)
             if isinstance(val, str):
                 if len(val) > MAX_STR_LEN:
-                    val = '...{}'.format(val[-MAX_STR_LEN:])
+                    val = "...{}".format(val[-MAX_STR_LEN:])
                 val = '"{}"'.format(val)
             repr_str += "  {}={},\n".format(attr, val)
 
@@ -33,7 +33,7 @@ class Track(object):
             if val.__doc__ is None:
                 raise ValueError("{} has no documentation".format(prop))
 
-            val_type_str = val.__doc__.split(':')[0]
+            val_type_str = val.__doc__.split(":")[0]
             repr_str += "  {}: {},\n".format(prop, val_type_str)
 
         repr_str += ")"
@@ -41,3 +41,119 @@ class Track(object):
 
     def to_jams(self):
         raise NotImplementedError
+
+
+class MultiTrack(Track):
+    """MultiTrack class.
+
+    A multitrack class is a collection of track objects and their associated audio
+    that can be mixed together.
+    A multitrack is iteslf a Track, and can have its own associated audio (such as
+    a mastered mix), its own metadata and its own annotations.
+
+    Attributes:
+        tracks (dict): {track_id: Track}
+        track_audio_attribute (str): the name of the attribute of Track which 
+            returns the audio to be mixed
+    """
+
+    def __init__(self, tracks, track_audio_attribute):
+        """Inits MultiTrack with tracks and audio attribute"""
+        self.tracks = tracks
+        self.track_audio_attribute = track_audio_attribute
+
+    def get_target(self, track_keys, weights=None, average=True, enforce_length=True):
+        """Get target which is a linear mixture of tracks
+
+        Args:
+            track_keys (list): list of track keys to mix together
+            weights (list or None): list of positive scalars to be used in the average
+            average (bool): if True, computes a weighted average of the tracks
+                if False, computes a weighted sum of the tracks
+            enforce_length (bool): If True, raises ValueError if the tracks are 
+                not the same length. If False, pads audio with zeros to match the length
+                of the longest track
+        
+        Returns:
+            target (np.ndarray): target audio with shape (n_channels, n_samples)
+
+        Raises:
+            ValueError: 
+                if sample rates of the tracks are not equal
+                if enforce_length=True and lengths are not equal
+
+        """
+        signals = []
+        lengths = []
+        sample_rates = []
+        for k in track_keys:
+            audio, sample_rate = getattr(self.tracks[k], self.track_audio_attribute)()
+            # ensure all signals are shape (n_channels, n_samples)
+            if len(audio.shape) == 1:
+                audio = audio[np.newaxis, :]
+            signals.append(audio)
+            lengths.append(audio.shape[1])
+            sample_rates.append(sample_rate)
+
+        if len(set(sample_rates)) > 1:
+            raise ValueError(
+                "Sample rates for tracks {} are not equal: {}".format(
+                    track_keys, sample_rates
+                )
+            )
+
+        max_length = np.max(lengths)
+        if any([l != max_length for l in lengths]):
+            if enforce_length:
+                raise ValueError(
+                    "Track's {} audio are not the same length {}. Use enforce_length=False to pad with zeros.".format(
+                        track_keys, lengths
+                    )
+                )
+            else:
+                # pad signals to the max length
+                signals = [
+                    np.pad(signal, ((0, 0), (0, max_length - signal.shape[1])))
+                    for signal in signals
+                ]
+
+        if weights is None:
+            weights = np.ones((len(track_keys),))
+
+        target = np.average(signals, axis=0, weights=weights)
+        if not average:
+            target *= np.sum(weights)
+
+        return target
+
+    def get_random_target(self, n_tracks=None, min_weight=0.3, max_weight=1.0):
+        """Get a random target by combining a random selection of tracks with random weights
+
+        Args:
+            n_tracks (int or None): number of tracks to randomly mix. If None, uses all tracks
+            min_weight (float): minimum possible weight when mixing
+            max_weight (float): maximum possible weight when mixing
+        
+        Returns:
+            target (np.ndarray): mixture audio with shape (n_samples, n_channels)
+            tracks (list): list of keys of included tracks
+            weights (list): list of weights used to mix tracks
+        """
+        tracks = list(self.tracks.keys())
+        if n_tracks is not None and n_tracks < len(tracks):
+            tracks = np.random.choice(tracks, n_tracks, replace=False)
+
+        weights = np.random.uniform(low=min_weight, high=max_weight, size=len(tracks))
+        target = self.get_target(tracks, weights=weights)
+        return target, tracks, weights
+
+    def get_mix(self):
+        """Create a linear mixture given a subset of tracks.
+
+        Args:
+            track_keys (list): list of track keys to mix together
+        
+        Returns:
+            target (np.ndarray): mixture audio with shape (n_samples, n_channels)
+        """
+        return self.get_target(list(self.tracks.keys()))
diff --git a/mirdata/version.py b/mirdata/version.py
index bc032fb..f3b35f0 100644
--- a/mirdata/version.py
+++ b/mirdata/version.py
@@ -2,5 +2,5 @@
 # -*- coding: utf-8 -*-
 """Version info"""
 
-short_version = '0.2'
-version = '0.2.0'
+short_version = "0.2"
+version = "0.2.1"


Test patch:
diff --git a/tests/test_track.py b/tests/test_track.py
index b15104e..a25c750 100644
--- a/tests/test_track.py
+++ b/tests/test_track.py
@@ -2,24 +2,25 @@
 
 import sys
 import pytest
+import numpy as np
 
 from mirdata import track
 
 if sys.version_info.major == 3:
-    builtin_module_name = 'builtins'
+    builtin_module_name = "builtins"
 else:
-    builtin_module_name = '__builtin__'
+    builtin_module_name = "__builtin__"
 
 
 def test_track_repr():
     class TestTrack(track.Track):
         def __init__(self):
-            self.a = 'asdf'
+            self.a = "asdf"
             self.b = 1.2345678
-            self.c = {1: 'a', 'b': 2}
-            self._d = 'hidden'
+            self.c = {1: "a", "b": 2}
+            self._d = "hidden"
             self.e = None
-            self.long = 'a' + 'b' * 50 + 'c' * 50
+            self.long = "a" + "b" * 50 + "c" * 50
 
         @property
         def f(self):
@@ -36,7 +37,7 @@ def test_track_repr():
 
     expected1 = """Track(\n  a="asdf",\n  b=1.2345678,\n  """
     expected2 = """c={1: 'a', 'b': 2},\n  e=None,\n  """
-    expected3 = """long="...{}",\n  """.format('b' * 50 + 'c' * 50)
+    expected3 = """long="...{}",\n  """.format("b" * 50 + "c" * 50)
     expected4 = """f: ThisObjectType,\n  g: I have an improper docstring,\n)"""
 
     test_track = TestTrack()
@@ -54,3 +55,160 @@ def test_track_repr():
     bad_track = NoDocsTrack()
     with pytest.raises(ValueError):
         bad_track.__repr__()
+
+
+def test_multitrack():
+    class TestTrack(track.Track):
+        def __init__(self, key):
+            self.key = key
+
+        def f(self):
+            return np.random.uniform(-1, 1, (2, 100)), 1000
+
+    track_keys = ["a", "b", "c"]
+    tracks = {k: TestTrack(k) for k in track_keys}
+
+    mtrack = track.MultiTrack(tracks, "f")
+
+    target1 = mtrack.get_target(["a", "c"])
+    assert target1.shape == (2, 100)
+    assert np.max(np.abs(target1)) <= 1
+
+    target2 = mtrack.get_target(["b", "c"], weights=[0.5, 0.2])
+    assert target2.shape == (2, 100)
+    assert np.max(np.abs(target2)) <= 1
+
+    target3 = mtrack.get_target(["b", "c"], weights=[0.5, 5])
+    assert target3.shape == (2, 100)
+    assert np.max(np.abs(target3)) <= 1
+
+    target4 = mtrack.get_target(["a", "c"], average=False)
+    assert target4.shape == (2, 100)
+    assert np.max(np.abs(target4)) <= 2
+
+    target5 = mtrack.get_target(["a", "c"], average=False, weights=[0.1, 0.5])
+    assert target5.shape == (2, 100)
+    assert np.max(np.abs(target5)) <= 0.6
+
+    random_target1, t1, w1 = mtrack.get_random_target(n_tracks=2)
+    assert random_target1.shape == (2, 100)
+    assert np.max(np.abs(random_target1)) <= 1
+    assert len(t1) == 2
+    assert len(w1) == 2
+    assert np.all(w1 >= 0.3)
+    assert np.all(w1 <= 1.0)
+
+    random_target2, t2, w2 = mtrack.get_random_target(n_tracks=5)
+    assert random_target2.shape == (2, 100)
+    assert np.max(np.abs(random_target2)) <= 1
+    assert len(t2) == 3
+    assert len(w2) == 3
+    assert np.all(w2 >= 0.3)
+    assert np.all(w2 <= 1.0)
+
+    random_target3, t3, w3 = mtrack.get_random_target()
+    assert random_target3.shape == (2, 100)
+    assert np.max(np.abs(random_target3)) <= 1
+    assert len(t3) == 3
+    assert len(w3) == 3
+    assert np.all(w3 >= 0.3)
+    assert np.all(w3 <= 1.0)
+
+    random_target4, t4, w4 = mtrack.get_random_target(
+        n_tracks=2, min_weight=0.1, max_weight=0.4
+    )
+    assert random_target4.shape == (2, 100)
+    assert np.max(np.abs(random_target4)) <= 1
+    assert len(t4) == 2
+    assert len(w4) == 2
+    assert np.all(w4 >= 0.1)
+    assert np.all(w4 <= 0.4)
+
+    mix = mtrack.get_mix()
+    assert mix.shape == (2, 100)
+
+
+def test_multitrack_unequal_len():
+    class TestTrack(track.Track):
+        def __init__(self, key):
+            self.key = key
+
+        def f(self):
+            return np.random.uniform(-1, 1, (2, np.random.randint(50, 100))), 1000
+
+    track_keys = ["a", "b", "c"]
+    tracks = {k: TestTrack(k) for k in track_keys}
+
+    mtrack = track.MultiTrack(tracks, "f")
+
+    with pytest.raises(ValueError):
+        mtrack.get_target(["a", "b", "c"])
+
+    target1 = mtrack.get_target(["a", "b", "c"], enforce_length=False)
+    assert target1.shape[0] == 2
+    assert np.max(np.abs(target1)) <= 1
+
+    target2 = mtrack.get_target(["a", "b", "c"], average=False, enforce_length=False)
+    assert target2.shape[0] == 2
+    assert np.max(np.abs(target2)) <= 3
+
+
+def test_multitrack_unequal_sr():
+    class TestTrack(track.Track):
+        def __init__(self, key):
+            self.key = key
+
+        def f(self):
+            return np.random.uniform(-1, 1, (2, 100)), np.random.randint(10, 1000)
+
+    track_keys = ["a", "b", "c"]
+    tracks = {k: TestTrack(k) for k in track_keys}
+
+    mtrack = track.MultiTrack(tracks, "f")
+
+    with pytest.raises(ValueError):
+        mtrack.get_target(["a", "b", "c"])
+
+
+def test_multitrack_mono():
+    ### no first channel - audio shapes (100,)
+    class TestTrack(track.Track):
+        def __init__(self, key):
+            self.key = key
+
+        def f(self):
+            return np.random.uniform(-1, 1, (100)), 1000
+
+    track_keys = ["a", "b", "c"]
+    tracks = {k: TestTrack(k) for k in track_keys}
+
+    mtrack = track.MultiTrack(tracks, "f")
+
+    target1 = mtrack.get_target(["a", "c"])
+    assert target1.shape == (1, 100)
+    assert np.max(np.abs(target1)) <= 1
+
+    target1 = mtrack.get_target(["a", "c"], average=False)
+    assert target1.shape == (1, 100)
+    assert np.max(np.abs(target1)) <= 2
+
+    ### one channel mono shape (1, 100)
+    class TestTrack1(track.Track):
+        def __init__(self, key):
+            self.key = key
+
+        def f(self):
+            return np.random.uniform(-1, 1, (1, 100)), 1000
+
+    track_keys = ["a", "b", "c"]
+    tracks = {k: TestTrack1(k) for k in track_keys}
+
+    mtrack = track.MultiTrack(tracks, "f")
+
+    target1 = mtrack.get_target(["a", "c"])
+    assert target1.shape == (1, 100)
+    assert np.max(np.abs(target1)) <= 1
+
+    target1 = mtrack.get_target(["a", "c"], average=False)
+    assert target1.shape == (1, 100)
+    assert np.max(np.abs(target1)) <= 2