add pre-initialization to signals objects

neurodsp/sim/multi.py

@@ -79,13 +79,13 @@ def sim_across_values(sim_func, sim_params, return_type='object'):
     >>> from neurodsp.sim.update import ParamIter
     >>> base_params = {'n_seconds' : 2, 'fs' : 250, 'exponent' : None}
     >>> param_iter = ParamIter(base_params, 'exponent', [-2, 1, 0])
-    >>> sims = sim_multi_across_values(sim_powerlaw, param_iter)
+    >>> sims = sim_across_values(sim_powerlaw, param_iter)
 
     Simulate multiple powerlaw signals from manually defined set of simulation parameters:
 
     >>> params = [{'n_seconds' : 2, 'fs' : 250, 'exponent' : -2},
     ...           {'n_seconds' : 2, 'fs' : 250, 'exponent' : -1}]
-    >>> sims = sim_multi_across_values(sim_powerlaw, params)
+    >>> sims = sim_across_values(sim_powerlaw, params)
     """
 
     base = get_base_params(sim_params)

neurodsp/sim/signals.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from neurodsp.utils.core import listify
+from neurodsp.utils.data import compute_nsamples
 from neurodsp.sim.params import get_base_params, drop_base_params, get_param_values
 
 ###################################################################################################
@@ -15,8 +16,9 @@ class Simulations():
 
     Parameters
     ----------
-    signals : 1d or 2nd array, optional
-        The simulated signals, organized as [n_sims, sig_length].
+    signals : 1d or 2nd array or int, optional
+        If array, the simulated signals, organized as [n_sims, sig_length].
+        If int, the number of expected simulations, used to pre-initialize array.
     params : dict, optional
         The simulation parameters that were used to create the simulations.
     sim_func : str or callable, optional
@@ -31,10 +33,17 @@ class Simulations():
     def __init__(self, signals=None, params=None, sim_func=None):
         """Initialize Simulations object."""
 
-        self.signals = np.atleast_2d(signals) if signals is not None else np.array([])
+        if signals is None:
+            signals = np.array([])
+        elif isinstance(signals, int):
+            n_samples = compute_nsamples(params['n_seconds'], params['fs'])
+            signals = np.zeros((signals, n_samples))
+        self.signals = np.atleast_2d(signals)
+
         self._base_params = None
         self._params = None
         self.add_params(params)
+
         self.sim_func = sim_func.__name__ if callable(sim_func) else sim_func
 
 
@@ -102,7 +111,7 @@ def add_params(self, params):
 
         Parameters
         ----------
-        params : dict, optional
+        params : dict or None
             The simulation parameter definition(s).
         """
 
@@ -111,13 +120,38 @@ def add_params(self, params):
             self._params = drop_base_params(params)
 
 
+    def add_signal(self, signal, index=None):
+        """Add a signal to the current object.
+
+        Parameters
+        ----------
+        signal : 1d array
+            A simulated signal to add to the object.
+        index : int
+            Index to insert the new signal in the signals attribute.
+        """
+
+        if index is not None:
+            self.signals[index, :] = signal
+        else:
+            if not self.signals.size:
+                self.signals = np.atleast_2d(signal)
+            else:
+                try:
+                    self.signals = np.vstack([self.signals, signal])
+                except ValueError as array_value_error:
+                    msg = 'Size of the added signal is not consistent with existing signals.'
+                    raise ValueError(msg) from array_value_error
+
+
 class VariableSimulations(Simulations):
     """Data object for a set of simulated signals with variable parameter definitions.
 
     Parameters
     ----------
-    signals : 2nd array, optional
-        The simulated signals, organized as [n_sims, sig_length].
+    signals : 2nd array or int, optional
+        If array, the simulated signals, organized as [n_sims, sig_length].
+        If int, the number of expected simulations, used to pre-initialize array.
     params : list of dict, optional
         The simulation parameters for each of the simulations.
     sim_func : str, optional
@@ -186,25 +220,30 @@ def add_params(self, params):
         if params:
 
             params = listify(params)
-            base_params = get_base_params(params[0])
-            cparams = [drop_base_params(el) for el in params]
 
-            if not self.has_params:
-                if len(self) > 1 and len(self) > len(cparams):
-                    msg = 'Cannot add parameters to object without existing parameter values.'
-                    raise ValueError(msg)
+            base_params = get_base_params(params[0])
+            if not self._base_params:
                 self._base_params = base_params
-                self._params = cparams
-
             else:
-                self._params.extend(cparams)
+                msg = 'Base params have to match existing parameters.'
+                assert base_params == self._base_params, msg
+
+            cparams = [drop_base_params(el) for el in params]
+            if cparams[0]:
+                if not self.has_params:
+                    if len(self) > 1 and len(self) > len(cparams):
+                        msg = 'Cannot add parameters to object without existing parameter values.'
+                        raise ValueError(msg)
+                    self._params = cparams
+                else:
+                    self._params.extend(cparams)
 
         else:
             if self.has_params:
                 raise ValueError('Must add parameters if object already has them.')
 
 
-    def add_signal(self, signal, params=None):
+    def add_signal(self, signal, params=None, index=None):
         """Add a signal to the current object.
 
         Parameters
@@ -215,16 +254,11 @@ def add_signal(self, signal, params=None):
             Parameter definition for the added signal.
             If current object does not include parameters, should be empty.
             If current object does include parameters, this input is required.
+        index : int
+            Index to insert the new signal in the signals attribute.
         """
 
-        if not self.signals.size:
-            self.signals = np.atleast_2d(signal)
-        else:
-            try:
-                self.signals = np.vstack([self.signals, signal])
-            except ValueError as array_value_error:
-                msg = 'Size of the added signal is not consistent with existing signals.'
-                raise ValueError(msg) from array_value_error
+        super().add_signal(signal, index=index)
         self.add_params(params)