Skip to content

Commit

Permalink
Merge pull request #88 from muon-spectroscopy-computational-project/8…
Browse files Browse the repository at this point in the history 
…6_apply_results_function

Do not reshape results_function evaluations #86
  • Loading branch information
@patrick-austin
patrick-austin authored Nov 7, 2023
2 parents 26df005 + 7d4f456 commit a0d87fe
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 50 deletions.
10 changes: 5 additions & 5 deletions muspinsim/experiment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -263,7 +263,6 @@ def dissipation_operators(self):
)

if self._dops is None:

# Create a copy of the system
sys = self._system.clone()

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
83 changes: 45 additions & 38 deletions muspinsim/input/keyword.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,6 @@ def _store_values(self, block):
self._values.append(b)

def evaluate(self, **variables):

allvars = {**variables, **self._constants}

def expreval(expr):
Expand All @@ -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:
Expand All @@ -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"
Expand Down Expand Up @@ -347,15 +360,13 @@ 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
default = "muspinsim"


class KWSpins(MuSpinKeyword):

name = "spins"
expr_size_bounds = (1, np.inf)
block_size = 1
Expand All @@ -364,7 +375,6 @@ class KWSpins(MuSpinKeyword):


class KWCelio(MuSpinKeyword):

name = "celio"
expr_size_bounds = (1, 2)
block_size = 1
Expand All @@ -373,7 +383,6 @@ class KWCelio(MuSpinKeyword):


class KWPolarization(MuSpinExpandKeyword):

name = "polarization"
expr_size_bounds = (1, 3)
block_size = 1
Expand All @@ -387,7 +396,6 @@ class KWPolarization(MuSpinExpandKeyword):


class KWField(MuSpinExpandKeyword):

name = "field"
block_size = 1
accept_range = True
Expand All @@ -397,7 +405,6 @@ class KWField(MuSpinExpandKeyword):


class KWIntrinsicField(MuSpinExpandKeyword):

name = "intrinsic_field"
block_size = 1
accept_range = True
Expand All @@ -407,7 +414,6 @@ class KWIntrinsicField(MuSpinExpandKeyword):


class KWTime(MuSpinExpandKeyword):

name = "time"
block_size = 1
accept_range = True
Expand All @@ -416,7 +422,6 @@ class KWTime(MuSpinExpandKeyword):


class KWXAxis(MuSpinKeyword):

name = "x_axis"
block_size = 1
accept_range = False
Expand All @@ -431,7 +436,6 @@ class KWXAxis(MuSpinKeyword):


class KWYAxis(MuSpinKeyword):

name = "y_axis"
block_size = 1
accept_range = False
Expand All @@ -444,7 +448,6 @@ class KWYAxis(MuSpinKeyword):


class KWAverageAxes(MuSpinKeyword):

name = "average_axes"
block_size = 1
accept_range = True
Expand All @@ -460,7 +463,6 @@ class KWAverageAxes(MuSpinKeyword):


class KWOrientation(MuSpinExpandKeyword):

name = "orientation"
expr_size_bounds = (1, 4)
block_size = 1
Expand All @@ -474,7 +476,6 @@ def _default_args(self, mode="zyz"):


class KWTemperature(MuSpinExpandKeyword):

name = "temperature"
block_size = 1
accept_range = True
Expand All @@ -484,7 +485,6 @@ class KWTemperature(MuSpinExpandKeyword):

# Couplings
class KWZeeman(MuSpinCouplingKeyword):

name = "zeeman"
expr_size_bounds = (3, 3)
block_size = 1
Expand All @@ -496,7 +496,6 @@ def _default_args(self, i):


class KWDipolar(MuSpinCouplingKeyword):

name = "dipolar"
block_size = 1
expr_size_bounds = (3, 3)
Expand All @@ -507,7 +506,6 @@ def _default_args(self, i, j):


class KWHyperfine(MuSpinCouplingKeyword):

name = "hyperfine"
block_size = 3
expr_size_bounds = (3, 3)
Expand All @@ -518,7 +516,6 @@ def _default_args(self, i, j=None):


class KWQuadrupolar(MuSpinCouplingKeyword):

name = "quadrupolar"
block_size = 3
expr_size_bounds = (3, 3)
Expand All @@ -531,7 +528,6 @@ def _default_args(self, i):


class KWDissipation(MuSpinCouplingKeyword):

name = "dissipation"
block_size = 1
expr_size_bounds = (1, 1)
Expand All @@ -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)
Expand All @@ -563,7 +558,6 @@ class KWFittingVariables(MuSpinKeyword):
]

def _store_values(self, block):

variables = list(self._constants.keys())

self._values = []
Expand All @@ -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)
Expand All @@ -599,7 +592,6 @@ class KWFittingData(MuSpinExpandKeyword):


class KWFittingMethod(MuSpinKeyword):

ACCEPTED_FITTING_METHODS = ["nelder-mead", "lbfgs", "least-squares"]

name = "fitting_method"
Expand All @@ -617,7 +609,6 @@ class KWFittingMethod(MuSpinKeyword):


class KWFittingTolerance(MuSpinKeyword):

name = "fitting_tolerance"
block_size = 1
accept_range = False
Expand All @@ -630,18 +621,34 @@ class KWFittingTolerance(MuSpinKeyword):


class KWResultsFunction(MuSpinExpandKeyword):

name = "results_function"
block_size = 1
accept_range = False
default = "y"
_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
Expand Down
Loading

0 comments on commit a0d87fe

Please sign in to comment.