Fix normalized errors in callback

brian2modelfitting/fitter.py

@@ -535,7 +535,9 @@ def fit(self, optimizer, metric=None, n_rounds=1, callback='text',
                     output_dim = DIMENSIONLESS
                 else:
                     output_dim = self.output_dim
-                error_dim = self.metric.get_dimensions(output_dim)
+                # Correct the units for the normalization factor
+                error_dim = self.metric.get_normalized_dimensions(output_dim)
+                best_error = Quantity(float(self.best_error), dim=error_dim)
                 errors = Quantity(errors, dim=error_dim)
                 param_dicts = [{p: Quantity(v, dim=self.model[p].dim)
                                 for p, v in zip(self.parameter_names,
@@ -545,13 +547,13 @@ def fit(self, optimizer, metric=None, n_rounds=1, callback='text',
                 param_dicts = [{p: v for p, v in zip(self.parameter_names,
                                                      one_param_set)}
                                for one_param_set in parameters]
+                best_error = self.best_error
 
             if callback(param_dicts,
                         errors,
                         self.best_params,
-                        self.best_error,
+                        best_error,
                         index) is True:
-                print('Stopping simulation')
                 break
 
         return self.best_params, self.best_error
@@ -916,7 +918,7 @@ def _callback_wrapper(params, iter, resid, *args, **kwds):
             error = mean(resid**2)
             errors.append(error)
             if self.use_units:
-                error_dim = self.output_dim**2
+                error_dim = self.output_dim**2 * get_dimensions(normalization)**2
                 all_errors = Quantity(errors, dim=error_dim)
                 params = {p: Quantity(val, dim=self.model[p].dim)
                           for p, val in params.items()}

brian2modelfitting/metric.py

@@ -7,7 +7,7 @@
 except ImportError:
     warnings.warn('eFEL package not found.')
 from itertools import repeat
-from brian2 import Hz, second, Quantity, ms, us
+from brian2 import Hz, second, Quantity, ms, us, get_dimensions
 from brian2.units.fundamentalunits import check_units, in_unit, DIMENSIONLESS
 from numpy import (array, sum, square, reshape, abs, amin, digitize,
                    rint, arange, atleast_2d, NaN, float64, split, shape,)
@@ -158,6 +158,23 @@ def get_dimensions(self, output_dim):
         """
         return DIMENSIONLESS
 
+    def get_normalized_dimensions(self, output_dim):
+        """
+        The physical dimensions of the normalized error. This will be
+        the same as the dimensions returned by `~.Metric.get_dimensions` if
+        the ``normalization`` is not used or set to a dimensionless value.
+
+        Parameters
+        ----------
+        output_dim : `.Dimension`
+            The dimensions of the output variable.
+
+        Returns
+        -------
+        dim : `.Dimension`
+            The physical dimensions of the normalized error.
+        """
+        return DIMENSIONLESS
 
     @abc.abstractmethod
     def get_features(self, model_results, target_results, dt):
@@ -249,8 +266,8 @@ def calc(self, model_traces, data_traces, dt):
             ``(n_samples, )``.
         """
         start_steps = int(round(self.t_start/dt))
-        features = self.get_features(model_traces[:, :, start_steps:] * self.normalization,
-                                     data_traces[:, start_steps:] * self.normalization,
+        features = self.get_features(model_traces[:, :, start_steps:] * float(self.normalization),
+                                     data_traces[:, start_steps:] * float(self.normalization),
                                      dt)
         errors = self.get_errors(features)
 
@@ -332,7 +349,7 @@ def calc(self, model_spikes, data_spikes, dt):
             model_spikes = relevant_model_spikes
             data_spikes = relevant_data_spikes
         features = self.get_features(model_spikes, data_spikes, dt)
-        errors = self.get_errors(features) * self.normalization
+        errors = self.get_errors(features) * float(self.normalization)
 
         return errors
 
@@ -381,6 +398,8 @@ def get_errors(self, features):
     def get_dimensions(self, output_dim):
         return output_dim**2
 
+    def get_normalized_dimensions(self, output_dim):
+        return output_dim**2 * get_dimensions(self.normalization)**2
 
 class FeatureMetric(TraceMetric):
     def __init__(self, stim_times, feat_list, weights=None, combine=None,
@@ -483,7 +502,7 @@ def get_errors(self, features):
                 sample_error += total
             errors.append(sample_error)
 
-        return array(errors) * self.normalization
+        return array(errors) * float(self.normalization)
 
 
 class GammaFactor(SpikeMetric):
@@ -535,7 +554,7 @@ def get_features(self, traces, output, dt):
                                       rate_correction=self.rate_correction)
                 gf_for_sample.append(gf)
             all_gf.append(gf_for_sample)
-        return array(all_gf) * self.normalization
+        return array(all_gf) * float(self.normalization)
 
     def get_errors(self, features):
         errors = features.mean(axis=1)