Overwrite x_axis only during fitting #77

docs/docs/input.md

@@ -201,7 +201,7 @@ x_axis
     field
 ```
 
-Which range to use as the X axis of the simulation's output files. Must be another keyword that accepts a range, and the given keyword *must be specified as a range* in this input file. When fitting, this is also assumed to be the X axis of the data to fit, and the range specified for this keyword is overridden by the fitting data. By default it's `time`. 
+Which range to use as the X axis of the simulation's output files. Must be another keyword that accepts a range, and the given keyword *must be specified as a range* in this input file. When fitting, if x_axis is specified or the range for the default range `time` is set, then the user provided X values will not be used to perform the fit (as the fitting data's X values are used to optimise the `fitting_variables`). They will, however, be used to perform a final evaluation when writing the .dat file. This enables the fit to use all the available data, whilst also allowing the user to evaluate data for a particular region of interest, or at a different precision. If the range provided does not overlap with the experimental data, then the user will be warned that the they are extrapolating the result of the fit. By default it's `time`. 
 
 ### y_axis
 

muspinsim/__main__.py

@@ -140,7 +140,7 @@ def main(use_mpi=False):
             runner.config.save_output(name=None, path=out_path)
     else:
         fitter = FittingRunner(in_file)
-        fitter.run(name=None, path=out_path)
+        fitter.run()
 
         if mpi.is_root:
             rep_dname = inp_dir
@@ -157,6 +157,7 @@ def main(use_mpi=False):
                     rep_dname = inp_dir
 
             fitter.write_report(fname=rep_fname, path=rep_dname)
+            fitter.write_data(name=None, path=out_path)
 
     if mpi.is_root:
         t_end = datetime.now()

muspinsim/experiment.py

@@ -24,7 +24,12 @@ class ExperimentRunner:
     provide caching for any quantities that might not need to be recalculated
     between successive snapshots."""
 
-    def __init__(self, in_file: MuSpinInput, variables: dict = None):
+    def __init__(
+        self,
+        in_file: MuSpinInput,
+        variables: dict = None,
+        use_experimental_x_axis: bool = True,
+    ):
         """Set up an experiment as defined by a MuSpinInput object
 
         Prepare a set of calculations (for multiple files and averages) as
@@ -39,6 +44,11 @@ def __init__(self, in_file: MuSpinInput, variables: dict = None):
                                 running a fitting calculation, in which case
                                 results_function will not be applied when 'run' is
                                 called as it is already applied by the FittingRunner.
+            use_experimental_x_axis -- Only used for fittings. If True, then any x_axis
+                                       specified by the user will be overwritten (this
+                                       should be the case forfitting, but not for final
+                                       evaluation where the user may want a different
+                                       range), Default is True.
         """
         # Fix W0102:dangerous-default-value
         if variables is None:
@@ -52,7 +62,10 @@ def __init__(self, in_file: MuSpinInput, variables: dict = None):
             # values for simulation configurations. These are then broadcast
             # across all nodes, each of which runs its own slice of them, and
             # finally gathered back together
-            config = MuSpinConfig(in_file.evaluate(**variables))
+            config = MuSpinConfig(
+                params=in_file.evaluate(**variables),
+                use_experimental_x_axis=use_experimental_x_axis,
+            )
         else:
             config = MuSpinConfig()
 
@@ -263,7 +276,6 @@ def dissipation_operators(self):
             )
 
         if self._dops is None:
-
             # Create a copy of the system
             sys = self._system.clone()
 
@@ -298,7 +310,6 @@ def sparse_sum(sp_mat_list):
 
             self._dops = []
             for i, a in self._config.dissipation_terms.items():
-
                 op_x = sparse_sum(self._single_spinops[i, :, None] * x[:, None])
                 op_y = sparse_sum(self._single_spinops[i, :, None] * y[:, None])
                 op_p = SpinOperator(op_x + 1.0j * op_y, dim=self.system.dimension)

muspinsim/fitting.py

@@ -64,15 +64,14 @@ def __init__(self, inpfile: MuSpinInput):
         self._sol = None
         self._cached_results = None
 
-    def run(self, name=None, path="."):
+    def run(self):
         """Run a fitting calculation using Scipy, and returns the solution
 
         Returns:
             sol {scipy.OptimizeResult} -- The result of the optimisation.
         """
 
         if mpi.is_root:
-
             # Get the correct string for the method
             method = {
                 "nelder-mead": "nelder-mead",
@@ -116,9 +115,6 @@ def run(self, name=None, path="."):
             self._runner.config.results = self._runner.apply_results_function(
                 self._runner.config.results, dict(zip(self._xnames, self._sol["x"]))
             )
-
-            # And now save the last result
-            self._runner.config.save_output(name=name, path=path)
         else:
             while not self._done:
                 self._compute_result(self._x)
@@ -182,6 +178,31 @@ def _targfun_minimise(self, x):
             err = np.average(np.abs(y - self._ytarg))
             return err
 
+    def write_data(self, name: str = None, path: str = "."):
+        """Write the .dat file, using the x_axis given by the user or if not
+        provided, use the experimental data x values as during the fitting.
+
+        Arguments:
+            name {str} -- Root name to use for the files
+            path {str} -- Folder path to save the files in
+        """
+        vardict = dict(zip(self._xnames, self._x))
+        runner = ExperimentRunner(
+            self._input, variables=vardict, use_experimental_x_axis=False
+        )
+        if (
+            runner.config.x_axis != self._runner.config.x_axis
+            or runner.config.x_axis_values != self._runner.config.x_axis_values
+        ):
+            # If our x_axis differs from the one used in the fitting, use the
+            # user specified values for evaluation
+            y = runner.run()
+            runner.apply_results_function(y, vardict)
+            runner.config.save_output(name=name, path=path)
+        else:
+            # Otherwise, can use the values from the final fitting runner
+            self._runner.config.save_output(name=name, path=path)
+
     def write_report(self, fname=None, path="./"):
         """Write a report file with the contents of the fitting optimisation.
 
@@ -196,7 +217,6 @@ def write_report(self, fname=None, path="./"):
         """
 
         if mpi.is_root:
-
             # Final variable values
             variables = dict(zip(self._xnames, self._sol["x"]))
             config = MuSpinConfig(self._input.evaluate(**variables))

muspinsim/input/input.py

@@ -11,6 +11,7 @@
 
 from muspinsim.input.keyword import (
     InputKeywords,
+    KWXAxis,
     MuSpinEvaluateKeyword,
     MuSpinCouplingKeyword,
 )
@@ -66,7 +67,6 @@ def _make_blocks(file_stream):
     indent = None
 
     for i, l in enumerate(lines):
-
         # Remove any comments
         l = l.split("#", 1)[0]
 
@@ -119,10 +119,10 @@ def __init__(self, file_stream=None):
             "function": None,
             # When true indicates all fitting can be done after the simulation
             "single_simulation": True,
+            "evaluate_x_axis": False,
         }
 
         if file_stream is not None:
-
             raw_blocks, block_line_nums = _make_blocks(file_stream)
 
             # if we find errors when parsing fitting variables, we post an error
@@ -228,7 +228,6 @@ def evaluate(self, **variables):
         result = {"couplings": {}, "fitting_info": self.fitting_info}
 
         for name, KWClass in InputKeywords.items():
-
             if issubclass(KWClass, MuSpinCouplingKeyword):
                 if name in self._keywords:
                     for kwid, kw in self._keywords[name].items():
@@ -255,6 +254,12 @@ def evaluate(self, **variables):
                         result[name] = KWClass(variables=variables)
                 elif name in self._keywords:
                     kw = self._keywords[name]
+                    if kw.name == "x_axis" or kw.name == KWXAxis.default:
+                        # We perform the fitting on the experimental data, but
+                        # if the user has tried to specify x_axis, we should
+                        # evaluate on that for the output .dat file
+                        result["fitting_info"]["evaluate_x_axis"] = True
+
                     v = variables if issubclass(KWClass, MuSpinEvaluateKeyword) else {}
                     val = kw.evaluate(**v)
 

muspinsim/simconfig.py

@@ -79,7 +79,7 @@ class MuSpinConfig:
     """A class to store a configuration for a MuSpinSim simulation, including
     any ranges of parameters as determined by the input file."""
 
-    def __init__(self, params={}):
+    def __init__(self, params={}, use_experimental_x_axis: bool = True):
         """Initialise a MuSpinConfig object
 
         Initialise a MuSpinConfig object from values produced by the
@@ -88,6 +88,9 @@ def __init__(self, params={}):
         Arguments:
             params {dict} -- Dictionary of parameters as returned by
                              MuSpinInput.evaluate
+            use_experimental_x_axis {bool} -- Only used for fittings. If True, then any
+                                              x_axis specified by the user will be
+                                              overwritten. Default is True.
 
         """
 
@@ -129,7 +132,6 @@ def __init__(self, params={}):
 
         # Now inspect all parameters
         for iname, cname in _CDICT.items():
-
             try:
                 p = params[iname]
             except KeyError as exc:
@@ -180,26 +182,55 @@ def __init__(self, params={}):
         if finfo["fit"]:
             if len(self._file_ranges) > 0:
                 raise MuSpinConfigError("Can not have file ranges when fitting")
-            logging.warning(
-                "'x_axis' values will be overwritten by the experimental data"
-            )
-            # The x axis is overridden, whatever it is
-            xname = list(self._x_range.keys())[0]
+
+            xname, xvalue = list(self._x_range.items())[0]
             self._constants.pop(xname, None)  # Just in case it was here
 
-            # Special cases where loaded data will be a single value, but we
-            # actually need to load multiple
-            if xname in ["B", "intrinsic_B"]:
-                # Default to z axis (the same as in _validate_B below -
-                # the orientation is applied later)
-                self._x_range[xname] = np.array(
-                    [[0, 0, v] for v in finfo["data"][:, 0]]
-                )
-            else:
-                self._x_range[xname] = finfo["data"][:, 0]
-            if xname == "t":
-                # Special case
-                self._time_N = len(self._x_range[xname])
+            if xvalue is None:
+                # Can happen when xname keyword is a single value, and so not assigned
+                # In this case, we must use the experimental data, as that's the only
+                # x_range we have
+                use_experimental_x_axis = True
+                finfo["evaluate_x_axis"] = False
+
+            if use_experimental_x_axis:
+                # Special cases where loaded data will be a single value, but we
+                # actually need to load multiple
+                if xname in ["B", "intrinsic_B"]:
+                    # Default to z axis (the same as in _validate_B below -
+                    # the orientation is applied later)
+                    experimental_x_range = np.array(
+                        [[0, 0, v] for v in finfo["data"][:, 0]]
+                    )
+                else:
+                    experimental_x_range = finfo["data"][:, 0]
+
+                if finfo["evaluate_x_axis"]:
+                    logging.info(
+                        "Fitting will be performed on experimental data-points, "
+                        "specified 'x_axis' will only be used to generate final "
+                        ".dat file"
+                    )
+                    below_min = xvalue < experimental_x_range[0]
+                    above_max = xvalue > experimental_x_range[-1]
+                    if np.all(above_max) or np.all(below_min):
+                        logging.warning(
+                            "All points on specified 'x_axis' are outside the "
+                            "experimental range, fitted parameters may not be "
+                            "appropriate for extrapolation"
+                        )
+                    elif np.any(above_max) or np.any(below_min):
+                        logging.info(
+                            "Some points on specified 'x_axis' are outside the "
+                            "experimental range, fitted parameters may not be "
+                            "appropriate for extrapolation"
+                        )
+
+                self._x_range[xname] = experimental_x_range
+
+                if xname == "t":
+                    # Special case
+                    self._time_N = len(self._x_range[xname])
 
         # Check that a X axis was found
         if list(self._x_range.values())[0] is None:
@@ -464,7 +495,6 @@ def __len__(self):
         return len(self._configurations)
 
     def __getitem__(self, i):
-
         isint = isinstance(i, int)
         if isint:
             i = slice(i, i + 1)
@@ -473,8 +503,7 @@ def __getitem__(self, i):
 
         ans = []
 
-        for (fc, ac, xc) in self._configurations[i]:
-
+        for fc, ac, xc in self._configurations[i]:
             fd = _elems_from_arrayodict(fc, self._file_ranges)
             ad = _elems_from_arrayodict(ac, self._avg_ranges)
             xd = _elems_from_arrayodict(xc, self._x_range)
@@ -495,7 +524,6 @@ def _validate_name(self, v):
         return v[0]
 
     def _validate_spins(self, v):
-
         # Find isotopes
         isore = re.compile("([0-9]+)([A-Z][a-z]*|e)")
         m = isore.match(v)
@@ -541,7 +569,6 @@ def _validate_t(self, v):
         return v[0]
 
     def _validate_B(self, v):
-
         if len(v) == 1:
             v = np.array([0, 0, v[0]])  # The default direction is Z
         elif len(v) != 3:
@@ -550,7 +577,6 @@ def _validate_B(self, v):
         return v
 
     def _validate_intrinsic_B(self, v):
-
         if len(v) == 1:
             v = np.array([0, 0, v[0]])  # The default direction is Z
         elif len(v) != 3:
@@ -568,7 +594,6 @@ def _validate_mupol(self, v):
         return v
 
     def _validate_orient(self, v, mode):
-
         q = None
         w = 1.0  # Weight
         if len(v) == 2: