From 7d4f456a1b54bb55db505541fb36548e5c5524a3 Mon Sep 17 00:00:00 2001 From: Patrick Austin Date: Mon, 19 Jun 2023 14:03:40 +0000 Subject: [PATCH] Do not reshape results_function evaluations #86 --- muspinsim/experiment.py | 10 ++-- muspinsim/input/keyword.py | 83 ++++++++++++++++-------------- muspinsim/tests/test_experiment.py | 31 ++++++++--- 3 files changed, 74 insertions(+), 50 deletions(-) diff --git a/muspinsim/experiment.py b/muspinsim/experiment.py index d891468..100e7f4 100644 --- a/muspinsim/experiment.py +++ b/muspinsim/experiment.py @@ -263,7 +263,6 @@ def dissipation_operators(self): ) if self._dops is None: - # Create a copy of the system sys = self._system.clone() @@ -298,7 +297,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) @@ -334,13 +332,15 @@ def p_operator(self): return self._system.muon_operator(self.p) def apply_results_function(self, results: ArrayLike, variables: dict): - # We expect muspinsim to output arrays with shape N here - # but the evaluation will return an array with shape (1, N) instead + # We expect muspinsim to output arrays with shape (M_1, M_2, ... N) + # here but the evaluation will return an array with shape + # (1, M_1, M_2, ... N) instead, where N denotes the x_axis and M any + # file ranges that are in use return np.array( self._config.results_function.evaluate( **variables, x=self._config.x_axis_values, y=results ) - ).reshape(len(results)) + ).reshape(results.shape) def run(self): """Run the experiment diff --git a/muspinsim/input/keyword.py b/muspinsim/input/keyword.py index 4c107e0..5920344 100644 --- a/muspinsim/input/keyword.py +++ b/muspinsim/input/keyword.py @@ -271,7 +271,6 @@ def _store_values(self, block): self._values.append(b) def evaluate(self, **variables): - allvars = {**variables, **self._constants} def expreval(expr): @@ -288,8 +287,33 @@ class MuSpinExpandKeyword(MuSpinEvaluateKeyword): block_size_bounds = (1, 1) _functions = {**_math_functions, "range": _range} - def evaluate(self, **variables): + def _reshape_value( + self, evaluated_line: np.ndarray, line_length: int + ) -> "list[np.ndarray]": + """Reshape evaluated values so they can be appended to the list of all + values. + + Arguments: + evaluated_line {np.ndarray} -- Results of evaluating a line + representing LarkExpression(s) + line_length {int} -- Length of the LarkExpression(s), which may be + different from the dimensions of the evaluated + values + + Returns: + {list[np.ndarray]} -- The evaluated line, as a list of ndarrays + """ + if len(evaluated_line.shape) == 1: + return [evaluated_line] + elif line_length == 1: + if len(evaluated_line.shape) == 2: + return [evaluated_line[0]] + elif len(evaluated_line.shape) == 3: + return list(evaluated_line[0]) + else: + raise RuntimeError(f"Unable to evaluate expression for keyword {self.name}") + def evaluate(self, **variables): allvars = {**variables, **self._constants} if self.block_size != 1: @@ -301,23 +325,12 @@ def expreval(expr): eval_values = [] for line in self._values: eval_line = np.array([expreval(expr) for expr in line]) - - if len(line) == 1 and len(eval_line.shape) == 3: - eval_values += list(eval_line[0]) - elif len(line) == 1 and len(eval_line.shape) == 2: - eval_values += [eval_line[0]] - elif len(eval_line.shape) == 1: - eval_values += [eval_line] - else: - raise RuntimeError( - f"Unable to evaluate expression for keyword {self.name}" - ) + eval_values += self._reshape_value(eval_line, len(line)) return eval_values class MuSpinCouplingKeyword(MuSpinEvaluateKeyword): - name = "coupling_keyword" block_size = 1 default = "0 0 0" @@ -347,7 +360,6 @@ def id(self): # Now on to defining the actual keywords that are admitted in input files class KWName(MuSpinKeyword): - name = "name" block_size = 1 accept_range = False @@ -355,7 +367,6 @@ class KWName(MuSpinKeyword): class KWSpins(MuSpinKeyword): - name = "spins" expr_size_bounds = (1, np.inf) block_size = 1 @@ -364,7 +375,6 @@ class KWSpins(MuSpinKeyword): class KWCelio(MuSpinKeyword): - name = "celio" expr_size_bounds = (1, 2) block_size = 1 @@ -373,7 +383,6 @@ class KWCelio(MuSpinKeyword): class KWPolarization(MuSpinExpandKeyword): - name = "polarization" expr_size_bounds = (1, 3) block_size = 1 @@ -387,7 +396,6 @@ class KWPolarization(MuSpinExpandKeyword): class KWField(MuSpinExpandKeyword): - name = "field" block_size = 1 accept_range = True @@ -397,7 +405,6 @@ class KWField(MuSpinExpandKeyword): class KWIntrinsicField(MuSpinExpandKeyword): - name = "intrinsic_field" block_size = 1 accept_range = True @@ -407,7 +414,6 @@ class KWIntrinsicField(MuSpinExpandKeyword): class KWTime(MuSpinExpandKeyword): - name = "time" block_size = 1 accept_range = True @@ -416,7 +422,6 @@ class KWTime(MuSpinExpandKeyword): class KWXAxis(MuSpinKeyword): - name = "x_axis" block_size = 1 accept_range = False @@ -431,7 +436,6 @@ class KWXAxis(MuSpinKeyword): class KWYAxis(MuSpinKeyword): - name = "y_axis" block_size = 1 accept_range = False @@ -444,7 +448,6 @@ class KWYAxis(MuSpinKeyword): class KWAverageAxes(MuSpinKeyword): - name = "average_axes" block_size = 1 accept_range = True @@ -460,7 +463,6 @@ class KWAverageAxes(MuSpinKeyword): class KWOrientation(MuSpinExpandKeyword): - name = "orientation" expr_size_bounds = (1, 4) block_size = 1 @@ -474,7 +476,6 @@ def _default_args(self, mode="zyz"): class KWTemperature(MuSpinExpandKeyword): - name = "temperature" block_size = 1 accept_range = True @@ -484,7 +485,6 @@ class KWTemperature(MuSpinExpandKeyword): # Couplings class KWZeeman(MuSpinCouplingKeyword): - name = "zeeman" expr_size_bounds = (3, 3) block_size = 1 @@ -496,7 +496,6 @@ def _default_args(self, i): class KWDipolar(MuSpinCouplingKeyword): - name = "dipolar" block_size = 1 expr_size_bounds = (3, 3) @@ -507,7 +506,6 @@ def _default_args(self, i, j): class KWHyperfine(MuSpinCouplingKeyword): - name = "hyperfine" block_size = 3 expr_size_bounds = (3, 3) @@ -518,7 +516,6 @@ def _default_args(self, i, j=None): class KWQuadrupolar(MuSpinCouplingKeyword): - name = "quadrupolar" block_size = 3 expr_size_bounds = (3, 3) @@ -531,7 +528,6 @@ def _default_args(self, i): class KWDissipation(MuSpinCouplingKeyword): - name = "dissipation" block_size = 1 expr_size_bounds = (1, 1) @@ -543,7 +539,6 @@ def _default_args(self, i): # Fitting variables. This is a special case class KWFittingVariables(MuSpinKeyword): - name = "fitting_variables" block_size = 1 expr_size_bounds = (1, np.inf) @@ -563,7 +558,6 @@ class KWFittingVariables(MuSpinKeyword): ] def _store_values(self, block): - variables = list(self._constants.keys()) self._values = [] @@ -588,7 +582,6 @@ def _store_values(self, block): # Other fitting-related keywords class KWFittingData(MuSpinExpandKeyword): - name = "fitting_data" block_size = 1 expr_size_bounds = (1, 2) @@ -599,7 +592,6 @@ class KWFittingData(MuSpinExpandKeyword): class KWFittingMethod(MuSpinKeyword): - ACCEPTED_FITTING_METHODS = ["nelder-mead", "lbfgs", "least-squares"] name = "fitting_method" @@ -617,7 +609,6 @@ class KWFittingMethod(MuSpinKeyword): class KWFittingTolerance(MuSpinKeyword): - name = "fitting_tolerance" block_size = 1 accept_range = False @@ -630,7 +621,6 @@ class KWFittingTolerance(MuSpinKeyword): class KWResultsFunction(MuSpinExpandKeyword): - name = "results_function" block_size = 1 accept_range = False @@ -638,10 +628,27 @@ class KWResultsFunction(MuSpinExpandKeyword): _constants = {**_math_constants, **_phys_constants} _special_variables = ["x", "y"] + def _reshape_value( + self, evaluated_line: np.ndarray, line_length: int + ) -> "list[np.ndarray]": + """Reshaping should not be performed for results functions, which in + general may have any number of file ranges. This overrides and retains + the full dimensionality of the results, returning a list. + + Arguments: + evaluated_line {np.ndarray} -- Results of evaluating a line + representing LarkExpression(s) + line_length {int} -- Unused + + Returns: + {list[np.ndarray]} -- The first element from the evaluated line, + in a list + """ + return [evaluated_line[0]] + # Special configuration keyword class KWExperiment(MuSpinKeyword): - name = "experiment" block_size = 1 accept_range = False diff --git a/muspinsim/tests/test_experiment.py b/muspinsim/tests/test_experiment.py index 9003a98..8f558bd 100644 --- a/muspinsim/tests/test_experiment.py +++ b/muspinsim/tests/test_experiment.py @@ -11,7 +11,6 @@ class TestExperiment(unittest.TestCase): def test_create(self): - gmu = constants.MU_GAMMA ge = constants.ELEC_GAMMA @@ -73,7 +72,6 @@ def test_create(self): self.assertAlmostEqual(ertest.rho0.expectation(Sz_e), 0.0) def test_rho0(self): - stest = StringIO( """ spins @@ -117,7 +115,6 @@ def test_rho0(self): ) def test_run(self): - # Empty system stest = StringIO( """ @@ -205,7 +202,6 @@ def test_run(self): self.assertAlmostEqual(results[0], 0.5 / (1.0 + 4 * np.pi**2 * tau**2)) def test_run_results_function(self): - # Empty system, modifying range to be 0 to 1 stest = StringIO( """ @@ -270,6 +266,30 @@ def test_run_results_function(self): self.assertTrue(np.all(results == 0.5)) + def test_file_range(self): + """Test that we get the expected shape of results when using using file + ranges. + """ + stest = StringIO( + """ +spins + e mu +time + range(0, 10) +field + 0 + 10 +temperature + inf + 0 +""" + ) + itest = MuSpinInput(stest) + ertest = ExperimentRunner(itest) + + results = ertest.run() + self.assertTrue(results.shape == (2, 2, 100)) + def test_run_intrinsic_field(self): # Check results from rotating are different when using field or intrinsic_field stest = StringIO( @@ -318,7 +338,6 @@ def test_run_intrinsic_field(self): ) def test_run_celio(self): - # Empty system stest = StringIO( """ @@ -449,7 +468,6 @@ def test_run_celio(self): # For testing the faster time evolution method is run # correctly in the right scenarios def test_run_fast(self): - # System without muon first should revert to the # slower method stest = StringIO( @@ -561,7 +579,6 @@ def test_run_fast(self): self.assertFalse(ertest._T_inf_speedup) def test_dissipation(self): - # Simple system g = 1.0 stest = StringIO(