add to array function, add step functionality

stimuli/SinusAmplitudeModulation.py

@@ -1,5 +1,7 @@
 from stimuli.AbstractStimulus import AbstractStimulus
 import numpy as np
+from numba import jit, njit
+import time
 
 
 class SinusAmplitudeModulationStimulus(AbstractStimulus):
@@ -13,11 +15,13 @@ class SinusAmplitudeModulationStimulus(AbstractStimulus):
         self.duration = duration
 
     def value_at_time_in_s(self, time_point):
+        carrier = np.sin(2 * np.pi * self.carrier_frequency * time_point)
+
         if time_point < self.start_time or time_point > self.start_time + self.duration:
-            return 0
+            return self.amplitude * carrier
 
         am = (1 + self.contrast * np.sin(2*np.pi*self.modulation_frequency * time_point))
-        carrier = np.sin(2*np.pi*self.carrier_frequency*time_point)
+
         return self.amplitude * am * carrier
 
     def get_stimulus_start_ms(self):
@@ -28,3 +32,40 @@ class SinusAmplitudeModulationStimulus(AbstractStimulus):
 
     def get_amplitude(self):
         return self.contrast
+
+    def as_array(self, time_start, total_time, step_size):
+        carrier = self.carrier_frequency
+        amp = self.amplitude
+        mod_freq = self.modulation_frequency
+        contrast = self.contrast
+        start_time = self.start_time
+        duration = self.duration
+
+        values = convert_to_array(carrier, amp, mod_freq, contrast, start_time, duration, time_start, total_time, step_size)
+
+        return values
+
+
+@jit(nopython=True)
+def convert_to_array(carrier_freq, amplitude, modulation_freq, contrast, start_time, duration, time_start, total_time, step_size):
+    # if the whole stimulus time has the amplitude modulation just built it at once;
+    if time_start >= start_time and start_time+duration < time_start+total_time:
+        carrier = np.sin(2 * np.pi * carrier_freq * np.arange(start_time, total_time-start_time, step_size/1000))
+        modulation = 1 + contrast * np.sin(2 * np.pi * modulation_freq * np.arange(start_time, total_time-start_time, step_size/1000))
+        values = amplitude * carrier * modulation
+        return values
+
+    # if it is split into parts with and without amplitude modulation built it in parts:
+    values = np.array(0)
+
+    if time_start < start_time:
+        carrier_before_am = np.sin(2 * np.pi * carrier_freq * np.arange(time_start, start_time, step_size / 1000))
+        np.concatenate(values, amplitude * carrier_before_am)
+
+    if time_start < start_time+duration and time_start+total_time > start_time:
+        if time_start+total_time > start_time+duration:
+            carrier_during_am = np.sin(2 * np.pi * carrier_freq * np.arange(start_time, start_time+duration, step_size/1000))
+            am = 1 + contrast * np.sin(2 * np.pi * modulation_freq * np.arange(start_time, start_time+duration, step_size/1000))
+            np.concatenate(values, amplitude*am*carrier_during_am)
+
+    return values