From 027410d61d4a11fc55db1d2eccbb78c23369e971 Mon Sep 17 00:00:00 2001 From: roger Date: Wed, 24 Jun 2020 12:05:05 +0800 Subject: [PATCH 1/2] Adds support for returning stim only for certain trials --- DataProcessingTools/trialstructures.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/DataProcessingTools/trialstructures.py b/DataProcessingTools/trialstructures.py index 8098e5d..b86f642 100644 --- a/DataProcessingTools/trialstructures.py +++ b/DataProcessingTools/trialstructures.py @@ -122,12 +122,15 @@ def get_timestamps(self, event_label): return self.timestamps[idx], self.trialidx[idx], self.stimidx[idx] - def get_stim(self, stimidx=0): + def get_stim(self, stimidx=0, trialidx=None): """ Return the timestamp, identity and location of the stimulus at `stimidx` of every trial. """ - fidx = np.where(self.stimidx==stimidx)[0] + fidx = self.stimidx==stimidx + if trialidx is not None: + qidx = np.isin(self.trialidx, trialidx) + fidx = fidx & qidx p = re.compile("stimulus_on_([0-9]+)_([0-9]+)") location = [] identity = [] From 982491dd03adf3075cd2b47ec3b163def88bcb80 Mon Sep 17 00:00:00 2001 From: roger Date: Wed, 24 Jun 2020 12:05:20 +0800 Subject: [PATCH 2/2] Adds tests --- tests/test_trialstructures.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_trialstructures.py b/tests/test_trialstructures.py index f7c9a07..bdfa1e8 100644 --- a/tests/test_trialstructures.py +++ b/tests/test_trialstructures.py @@ -51,6 +51,11 @@ def test_working_memory(): assert location.count(2) == 431 assert location.count(3) == 312 + # only correct trials + reward_onset, ridx, _ = trials.get_timestamps("reward_on") + stim1_onset, identity, location = trials.get_stim(0, ridx) + assert len(stim1_onset) == 395 + testdir = "animal/20130923/session01" with DPT.misc.CWD(os.path.join(os.path.dirname(__file__),testdir)): trials = DPT.trialstructures.WorkingMemoryTrials()