Repository: NeuralEnsemble/python-neo

Issue:
AnalogSignal.rescale() does not preserve parent objects
The same is true for IrregularlySampledSignal and SpikeTrain.

Discussion:

Patch:
diff --git a/neo/core/analogsignal.py b/neo/core/analogsignal.py
index aa971c80..fec332ea 100644
--- a/neo/core/analogsignal.py
+++ b/neo/core/analogsignal.py
@@ -251,7 +251,7 @@ class AnalogSignal(BaseNeo, pq.Quantity):
         self.file_origin = getattr(obj, 'file_origin', None)
         self.description = getattr(obj, 'description', None)
 
-        # Parents objects
+        # Parent objects
         self.segment = getattr(obj, 'segment', None)
         self.channel_index = getattr(obj, 'channel_index', None)
 
@@ -437,7 +437,10 @@ class AnalogSignal(BaseNeo, pq.Quantity):
         new = self.__class__(signal=signal, units=to_u,
                              sampling_rate=self.sampling_rate)
         new._copy_data_complement(self)
+        new.channel_index = self.channel_index
+        new.segment = self.segment
         new.annotations.update(self.annotations)
+
         return new
 
     def duplicate_with_new_array(self, signal):
diff --git a/neo/core/container.py b/neo/core/container.py
index f59c1641..5f4535bf 100644
--- a/neo/core/container.py
+++ b/neo/core/container.py
@@ -3,8 +3,7 @@
 This module implements generic container base class that all neo container
 object inherit from.  It provides shared methods for all container types.
 
-:class:`Container` is derived from :class:`BaseNeo` but is
-defined in :module:`neo.core.analogsignalarray`.
+:class:`Container` is derived from :class:`BaseNeo`
 """
 
 # needed for python 3 compatibility
diff --git a/neo/core/irregularlysampledsignal.py b/neo/core/irregularlysampledsignal.py
index 62e9c4ac..8d919a8d 100644
--- a/neo/core/irregularlysampledsignal.py
+++ b/neo/core/irregularlysampledsignal.py
@@ -208,6 +208,10 @@ class IrregularlySampledSignal(BaseNeo, pq.Quantity):
         self.file_origin = getattr(obj, 'file_origin', None)
         self.description = getattr(obj, 'description', None)
 
+        # Parent objects
+        self.segment = getattr(obj, 'segment', None)
+        self.channel_index = getattr(obj, 'channel_index', None)
+
     def __repr__(self):
         '''
         Returns a string representing the :class:`IrregularlySampledSignal`.
@@ -456,6 +460,8 @@ class IrregularlySampledSignal(BaseNeo, pq.Quantity):
             signal = cf * self.magnitude
         new = self.__class__(times=self.times, signal=signal, units=to_u)
         new._copy_data_complement(self)
+        new.channel_index = self.channel_index
+        new.segment = self.segment
         new.annotations.update(self.annotations)
         return new
 
diff --git a/neo/core/spiketrain.py b/neo/core/spiketrain.py
index 5d874d95..79e3527d 100644
--- a/neo/core/spiketrain.py
+++ b/neo/core/spiketrain.py
@@ -330,12 +330,15 @@ class SpikeTrain(BaseNeo, pq.Quantity):
         if self.dimensionality == pq.quantity.validate_dimensionality(units):
             return self.copy()
         spikes = self.view(pq.Quantity)
-        return SpikeTrain(times=spikes, t_stop=self.t_stop, units=units,
-                          sampling_rate=self.sampling_rate,
-                          t_start=self.t_start, waveforms=self.waveforms,
-                          left_sweep=self.left_sweep, name=self.name,
-                          file_origin=self.file_origin,
-                          description=self.description, **self.annotations)
+        obj = SpikeTrain(times=spikes, t_stop=self.t_stop, units=units,
+                         sampling_rate=self.sampling_rate,
+                         t_start=self.t_start, waveforms=self.waveforms,
+                         left_sweep=self.left_sweep, name=self.name,
+                         file_origin=self.file_origin,
+                         description=self.description, **self.annotations)
+        obj.segment = self.segment
+        obj.unit = self.unit
+        return obj
 
     def __reduce__(self):
         '''


Test patch:
diff --git a/neo/test/coretest/test_analogsignal.py b/neo/test/coretest/test_analogsignal.py
index e56e48ab..bfadb827 100644
--- a/neo/test/coretest/test_analogsignal.py
+++ b/neo/test/coretest/test_analogsignal.py
@@ -302,7 +302,7 @@ class TestAnalogSignalArrayMethods(unittest.TestCase):
         self.signal1 = AnalogSignal(self.data1quant, sampling_rate=1*pq.kHz,
                                          name='spam', description='eggs',
                                          file_origin='testfile.txt', arg1='test')
-        self.signal1.segment = 1
+        self.signal1.segment = Segment()
         self.signal1.channel_index = ChannelIndex(index=[0])
 
     def test__compliant(self):
@@ -392,8 +392,8 @@ class TestAnalogSignalArrayMethods(unittest.TestCase):
     def test__copy_should_let_access_to_parents_objects(self):
         ##copy
         result =  self.signal1.copy()
-        self.assertEqual(result.segment, self.signal1.segment)
-        self.assertEqual(result.channel_index, self.signal1.channel_index)
+        self.assertIs(result.segment, self.signal1.segment)
+        self.assertIs(result.channel_index, self.signal1.channel_index)
         ## deep copy (not fixed yet)
         #result = copy.deepcopy(self.signal1)
         #self.assertEqual(result.segment, self.signal1.segment)
@@ -452,6 +452,11 @@ class TestAnalogSignalArrayMethods(unittest.TestCase):
         assert_array_equal(result.magnitude, self.data1.reshape(-1, 1))
         assert_same_sub_schema(result, self.signal1)
 
+        self.assertIsInstance(result.channel_index, ChannelIndex)
+        self.assertIsInstance(result.segment, Segment)
+        self.assertIs(result.channel_index, self.signal1.channel_index)
+        self.assertIs(result.segment, self.signal1.segment)
+
     def test__rescale_new(self):
         result = self.signal1.copy()
         result = result.rescale(pq.pA)
@@ -466,6 +471,11 @@ class TestAnalogSignalArrayMethods(unittest.TestCase):
         self.assertEqual(result.units, 1*pq.pA)
         assert_arrays_almost_equal(np.array(result), self.data1.reshape(-1, 1)*1000., 1e-10)
 
+        self.assertIsInstance(result.channel_index, ChannelIndex)
+        self.assertIsInstance(result.segment, Segment)
+        self.assertIs(result.channel_index, self.signal1.channel_index)
+        self.assertIs(result.segment, self.signal1.segment)
+
     def test__rescale_new_incompatible_ValueError(self):
         self.assertRaises(ValueError, self.signal1.rescale, pq.mV)
 
diff --git a/neo/test/coretest/test_irregularysampledsignal.py b/neo/test/coretest/test_irregularysampledsignal.py
index 1d379f65..c21cf68e 100644
--- a/neo/test/coretest/test_irregularysampledsignal.py
+++ b/neo/test/coretest/test_irregularysampledsignal.py
@@ -19,7 +19,7 @@ else:
     HAVE_IPYTHON = True
 
 from neo.core.irregularlysampledsignal import IrregularlySampledSignal
-from neo.core import Segment
+from neo.core import Segment, ChannelIndex
 from neo.test.tools import (assert_arrays_almost_equal, assert_arrays_equal,
                             assert_neo_object_is_compliant,
                             assert_same_sub_schema)
@@ -271,6 +271,8 @@ class TestIrregularlySampledSignalArrayMethods(unittest.TestCase):
                                                 description='eggs',
                                                 file_origin='testfile.txt',
                                                 arg1='test')
+        self.signal1.segment = Segment()
+        self.signal1.channel_index = ChannelIndex([0])
 
     def test__compliant(self):
         assert_neo_object_is_compliant(self.signal1)
@@ -345,6 +347,11 @@ class TestIrregularlySampledSignalArrayMethods(unittest.TestCase):
         assert_array_equal(result.times, self.time1quant)
         assert_same_sub_schema(result, self.signal1)
 
+        self.assertIsInstance(result.channel_index, ChannelIndex)
+        self.assertIsInstance(result.segment, Segment)
+        self.assertIs(result.channel_index, self.signal1.channel_index)
+        self.assertIs(result.segment, self.signal1.segment)
+
     def test__rescale_new(self):
         result = self.signal1.copy()
         result = result.rescale(pq.uV)
@@ -360,6 +367,11 @@ class TestIrregularlySampledSignalArrayMethods(unittest.TestCase):
         assert_arrays_almost_equal(np.array(result), self.data1.reshape(-1, 1)*1000., 1e-10)
         assert_array_equal(result.times, self.time1quant)
 
+        self.assertIsInstance(result.channel_index, ChannelIndex)
+        self.assertIsInstance(result.segment, Segment)
+        self.assertIs(result.channel_index, self.signal1.channel_index)
+        self.assertIs(result.segment, self.signal1.segment)
+
     def test__rescale_new_incompatible_ValueError(self):
         self.assertRaises(ValueError, self.signal1.rescale, pq.nA)
 
@@ -550,6 +562,11 @@ class TestIrregularlySampledSignalArrayMethods(unittest.TestCase):
         self.assertIsInstance(sig_as_q, pq.Quantity)
         assert_array_equal(self.data1, sig_as_q.magnitude.flat)
 
+    def test__copy_should_preserve_parent_objects(self):
+        result = self.signal1.copy()
+        self.assertIs(result.segment, self.signal1.segment)
+        self.assertIs(result.channel_index, self.signal1.channel_index)
+
 
 class TestIrregularlySampledSignalCombination(unittest.TestCase):
     def setUp(self):
diff --git a/neo/test/coretest/test_spiketrain.py b/neo/test/coretest/test_spiketrain.py
index 47c8dd23..4a1bbacd 100644
--- a/neo/test/coretest/test_spiketrain.py
+++ b/neo/test/coretest/test_spiketrain.py
@@ -1394,11 +1394,16 @@ class TestChanging(unittest.TestCase):
     def test__rescale(self):
         data = [3, 4, 5] * pq.ms
         train = SpikeTrain(data, t_start=0.5, t_stop=10.0)
+        train.segment = Segment()
+        train.unit = Unit()
         result = train.rescale(pq.s)
         assert_neo_object_is_compliant(train)
         assert_neo_object_is_compliant(result)
         assert_arrays_equal(train, result)
         self.assertEqual(result.units, 1 * pq.s)
+        self.assertIs(result.segment, train.segment)
+        self.assertIs(result.unit, train.unit)
+
 
     def test__rescale_same_units(self):
         data = [3, 4, 5] * pq.ms