redback 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

redback/simulate_transients.py

@@ -14,11 +14,11 @@ datadir = os.path.join(os.path.dirname(redback.__file__), 'tables')
 
 class SimulateGenericTransient(object):
     def __init__(self, model, parameters, times, model_kwargs, data_points,
-                 seed=1234, multiwavelength_transient=False, noise_term=0.2):
+                 seed=1234, multiwavelength_transient=False, noise_term=0.2, noise_type='gaussianmodel', extra_scatter=0.0):
         """
         A generic interface to simulating transients
 
-        :param model: String corresponding to redback model
+        :param model: String corresponding to redback model or a python function that can evaluate an SED.
         :param parameters: Dictionary of parameters describing a single transient
         :param times: Time values that the model is evaluated from
         :param model_kwargs: Additional keyword arguments, must include all the keyword arguments required by the model.
@@ -31,9 +31,17 @@ class SimulateGenericTransient(object):
                 and the data points are sampled in bands/frequency as well,
                 rather than just corresponding to one wavelength/filter.
                 This also allows the same time value to be sampled multiple times.
-        :param noise_term: Float. Factor which is multiplied by the model flux/magnitude to give the sigma.
-        """
-        self.model = redback.model_library.all_models_dict[model]
+        :param noise_type: String. Type of noise to add to the model.
+            Default is 'gaussianmodel' where sigma is noise_term * model.
+            Another option is 'gaussian' i.e., a simple Gaussian noise with sigma = noise_term.
+        :param noise_term: Float. Factor which is multiplied by the model flux/magnitude to give the sigma
+            or is sigma itself for 'gaussian' noise. Or the SNR for 'SNRbased' noise.
+        :param extra_scatter: Float. Sigma of normal added to output for additional scatter.
+        """
+        if model in redback.model_library.all_models_dict:
+            self.model = redback.model_library.all_models_dict[model]
+        else:
+            self.model = model 
         self.parameters = parameters
         self.all_times = times
         self.model_kwargs = model_kwargs
@@ -44,20 +52,25 @@ class SimulateGenericTransient(object):
         self.noise_term = noise_term
         random.seed(self.seed)
 
-        if multiwavelength_transient:
-            self.all_bands = self.model_kwargs.get('bands', None)
-            self.all_frequency = self.model_kwargs.get('frequency', None)
-            if self.all_bands is None and self.all_frequency is None:
-                raise ValueError('Must supply either bands or frequency to sample data points for an optical transient')
-            if self.all_bands is not None and self.all_frequency is None:
-                self.subset_bands = np.array(random.choices(self.all_bands, k=self.data_points))
-            if self.all_bands is None and self.all_frequency is not None:
-                self.subset_frequency = np.array(random.choices(self.all_frequency, k=self.data_points))
-            self.replacement = True
-            # allow times to be chosen repeatedly
+        self.all_bands = self.model_kwargs.get('bands', None)
+        self.all_frequency = self.model_kwargs.get('frequency', None)
+        if self.all_bands is None and self.all_frequency is None:
+            raise ValueError('Must supply either bands or frequency to sample data points for an optical transient')
         else:
-            # allow times to be chosen only once.
-            self.replacement = False
+            if multiwavelength_transient:
+                if self.all_bands is not None and self.all_frequency is None:
+                    self.subset_bands = np.array(random.choices(self.all_bands, k=self.data_points))
+                if self.all_bands is None and self.all_frequency is not None:
+                    self.subset_frequency = np.array(random.choices(self.all_frequency, k=self.data_points))
+                self.replacement = True
+                # allow times to be chosen repeatedly
+            else:
+                if self.all_bands is not None and self.all_frequency is None:
+                    self.subset_bands = self.data_points * [self.all_bands]
+                if self.all_bands is None and self.all_frequency is not None:
+                    self.subset_frequency = np.ones(self.data_points) * self.all_frequency
+                # allow times to be chosen only once.
+                self.replacement = False
         self.subset_times = np.sort(np.random.choice(self.all_times, size=self.data_points, replace=self.replacement))
 
         injection_kwargs = self.parameters.copy()
@@ -76,24 +89,47 @@ class SimulateGenericTransient(object):
         if 'frequency' in model_kwargs.keys():
             data['frequency'] = self.subset_frequency
         data['true_output'] = true_output
-        data['output'] = np.random.normal(true_output, self.noise_term * true_output)
-        data['output_error'] = self.noise_term * true_output
+
+        if noise_type == 'gaussianmodel':
+            noise = np.random.normal(0, self.noise_term * true_output, len(true_output))
+            output = true_output + noise
+            output_error = self.noise_term * true_output
+        elif noise_type == 'gaussian':
+            noise = np.random.normal(0, self.noise_term, len(true_output))
+            output = true_output + noise
+            output_error = self.noise_term
+        elif noise_type == 'SNRbased':
+            sigma = np.sqrt(true_output + np.min(true_output)/self.noise_term)
+            output_error = sigma
+            output = true_output + np.random.normal(0, sigma, len(true_output))
+        else:
+            logger.warning(f"noise_type {noise_type} not implemented.")
+            raise ValueError('noise_type must be either gaussianmodel, gaussian, or SNRBased')
+
+        if extra_scatter > 0:
+            extra_noise = np.random.normal(0, extra_scatter, len(true_output))
+            output = output + extra_noise
+            output_error = np.sqrt(output_error**2 + extra_noise**2)
+
+        data['output'] = output
+        data['output_error'] = output_error
         self.data = data
 
-        def save_transient(self, name):
-            """
-            Save the transient observations to a csv file.
-            This will save the full observational dataframe including non-detections etc.
-            This will save the data to a folder called 'simulated'
-            with the name of the transient and a csv file of the injection parameters
-
-            :param name: name to save transient.
-            """
-            bilby.utils.check_directory_exists_and_if_not_mkdir('simulated')
-            path = 'simulated/' + name + '.csv'
-            injection_path = 'simulated/' + name + '_injection_parameters.csv'
-            self.observations.to_csv(path, index=False)
-            self.parameters.to_csv(injection_path, index=False)
+    def save_transient(self, name):
+        """
+        Save the transient observations to a csv file.
+        This will save the full observational dataframe including non-detections etc.
+        This will save the data to a folder called 'simulated'
+        with the name of the transient and a csv file of the injection parameters
+
+        :param name: name to save transient.
+        """
+        bilby.utils.check_directory_exists_and_if_not_mkdir('simulated')
+        path = 'simulated/' + name + '.csv'
+        injection_path = 'simulated/' + name + '_injection_parameters.csv'
+        self.data.to_csv(path, index=False)
+        self.parameters=pd.DataFrame.from_dict([self.parameters])
+        self.parameters.to_csv(injection_path, index=False)
 
 class SimulateOpticalTransient(object):
     def __init__(self, model, parameters, pointings_database=None,
@@ -415,7 +451,7 @@ class SimulateOpticalTransient(object):
         Convert the circular field of view to a radius in radians.
         :return: survey_radius in radians
         """
-        survey_fov_sqrad = self.survey_fov_sqdeg*(np.pi/180.0)
+        survey_fov_sqrad = self.survey_fov_sqdeg*(np.pi/180.0)**2
         survey_radius = np.sqrt(survey_fov_sqrad/np.pi)
         # survey_radius = np.sqrt(self.survey_fov_sqdeg*((np.pi/180.0)**2.0)/np.pi)
         return survey_radius