diff --git a/doc/examples/profile/model.py b/doc/examples/profile/model.py index bd344219..773a730d 100644 --- a/doc/examples/profile/model.py +++ b/doc/examples/profile/model.py @@ -100,14 +100,14 @@ def mag(z, z1, z2, M1, M2, M3): # Set the fittable parameters. Note that the parameters to the function # after the first parameter *z* become fittable parameters. -sample["sin"].period.range(0, 100) -sample["sin"].phase.range(0, 1) sample["sin"].thickness.range(0, 1000) -sample["sin"].magnetism.M1.range(0, 10) -sample["sin"].magnetism.M2.range(0, 10) -sample["sin"].magnetism.M3.range(0, 10) -sample["sin"].magnetism.z1.range(0, 100) -sample["sin"].magnetism.z2.range(0, 100) +sample["sin"].profile_params["period"].range(0, 100) +sample["sin"].profile_params["phase"].range(0, 1) +sample["sin"].magnetism.profile_params["M1"].range(0, 10) +sample["sin"].magnetism.profile_params["M2"].range(0, 10) +sample["sin"].magnetism.profile_params["M3"].range(0, 10) +sample["sin"].magnetism.profile_params["z1"].range(0, 100) +sample["sin"].magnetism.profile_params["z2"].range(0, 100) # Define the model. Since this is a simulation, we need to define the # incident beam in terms of angles, wavelengths and dispersion. This diff --git a/refl1d/flayer.py b/refl1d/flayer.py index a3b7f46a..cae837b1 100644 --- a/refl1d/flayer.py +++ b/refl1d/flayer.py @@ -1,14 +1,19 @@ +from dataclasses import dataclass +import inspect +import numbers +from typing import Dict, Literal, Optional, Tuple import inspect from numpy import real, imag, asarray, broadcast_to -from bumps.parameter import Parameter, to_dict, Calculation +from bumps.parameter import Parameter, Calculation, Constant, Expression from refl1d.material import SLD -from refl1d.model import Layer +from refl1d.model import Layer, Stack from refl1d.magnetism import BaseMagnetism, Magnetism, DEFAULT_THETA_M from refl1d import util +@dataclass class FunctionalProfile(Layer): """ Generic profile function @@ -54,53 +59,111 @@ def linear(z, rhoL, rhoR): sample = L1 | profile | L3 """ + name: str + thickness: Parameter + interface: Constant + profile_source: str = None + profile_params: Dict[str, Parameter] = None + tol: float = 0 + magnetism: Optional[BaseMagnetism] + rho_start: Parameter = None + rho_end: Parameter = None + irho_start: Parameter = None + irho_end: Parameter = None + RESERVED = ("thickness", "interface", "profile", "tol", "magnetism", "name") # TODO: test that thickness(z) matches the thickness of the layer - def __init__(self, thickness=0, interface=0, profile=None, tol=1e-3, magnetism=None, name=None, **kw): + def __init__( + self, + thickness=0, + interface=0, + profile=None, + profile_source=None, + tol=1e-3, + magnetism=None, + name=None, + profile_params=None, + rho_start=None, + rho_end=None, + irho_start=None, + irho_end=None, + **kw, + ): if not name: name = profile.__name__ - if interface != 0: + if float(interface) != 0: raise NotImplementedError("interface not yet supported") - if profile is None: - raise TypeError("Need profile") + + if profile_source is not None: + self.profile_source = profile_source + import numpy as np + import scipy + + ctx = {"np": np, "scipy": scipy} + output = {} + exec(profile_source, ctx, output) + # assume that the function definition is the only thing in the source + profile_name = list(output.keys())[0] + self.profile = output[profile_name] + elif profile is not None: + self.profile = profile + self.profile_source = inspect.getsource(profile) + profile_name = profile.__name__ + else: + raise TypeError("Need profile or profile_source") + if not name: + name = profile_name + self.name = name self.thickness = Parameter.default(thickness, name=name + " thickness") - self.interface = Parameter.default(interface, name=name + " interface") - self.profile = profile + self.interface = Constant(interface, name=name + " interface") + self.rho_start = Parameter.default(rho_start, name=name + " rho_start") + self.rho_end = Parameter.default(rho_end, name=name + " rho_end") + self.irho_start = Parameter.default(irho_start, name=name + " irho_start") + self.irho_end = Parameter.default(irho_end, name=name + " irho_end") self.tol = tol self.magnetism = magnetism + if profile_params is None: + # use kw arguments to set initial values for profile parameters + self.profile_params = _init_profile_params(name, self.profile, kw, self.RESERVED) + else: + # no need to filter out reserved keywords + self.profile_params = _init_profile_params(name, self.profile, profile_params, ()) + # TODO: maybe make these lazy (and for magnetism below as well) - rho_start = _LayerLimit(self, isend=False, isrho=True) - irho_start = _LayerLimit(self, isend=False, isrho=False) - rho_end = _LayerLimit(self, isend=True, isrho=True) - irho_end = _LayerLimit(self, isend=True, isrho=False) - self.start = SLD(name + " start", rho=rho_start, irho=irho_start) - self.end = SLD(name + " end", rho=rho_end, irho=irho_end) + self._set_ends() + + def _set_ends(self): + rho_start = Calculation(self.name + " rho_start") + rho_start.set_function(lambda: real(self.profile(asarray([0.0]), **_get_values(self))[0])) + irho_start = Calculation(self.name + " irho_start") + irho_start.set_function(lambda: imag(self.profile(asarray([0.0]), **_get_values(self))[0])) + rho_end = Calculation(self.name + " rho_end") + rho_end.set_function(lambda: real(self.profile(asarray([self.thickness.value]), **_get_values(self))[0])) + irho_end = Calculation(self.name + " irho_end") + irho_end.set_function(lambda: imag(self.profile(asarray([self.thickness.value]), **_get_values(self))[0])) + self.rho_start.slot = rho_start + self.irho_start.slot = irho_start + self.rho_end.slot = rho_end + self.irho_end.slot = irho_end - _set_parameters(self, name, profile, kw, self.RESERVED) + @property + def start(self): + return SLD(self.name + " start", rho=self.rho_start, irho=self.irho_start) + + @property + def end(self): + return SLD(self.name + " end", rho=self.rho_end, irho=self.irho_end) def parameters(self): - P = _get_parameters(self) + P = {} + P.update(self.profile_params) P["thickness"] = self.thickness # P['interface'] = self.interface return P - def to_dict(self): - return to_dict( - { - "type": type(self).__name__, - "name": self.name, - "thickness": self.thickness, - "interface": self.interface, - "profile": self.profile, - "parameters": _get_parameters(self), - "tol": self.tol, - "magnetism": self.magnetism, - } - ) - def render(self, probe, slabs): Pw, Pz = slabs.microslabs(self.thickness.value) if len(Pw) == 0: @@ -117,6 +180,7 @@ def render(self, probe, slabs): # slabs.interface(self.interface.value) +@dataclass class FunctionalMagnetism(BaseMagnetism): """ Functional magnetism profile. @@ -138,31 +202,122 @@ class FunctionalMagnetism(BaseMagnetism): function. """ + name: str + profile_source: str = None + tol: float = 1e-3 + rhoM_start: Parameter = None # Calculation objects... + thetaM_start: Parameter = None + rhoM_end: Parameter = None + thetaM_end: Parameter = None + thickness: Parameter = None + profile_params: Dict[str, Parameter] = None + RESERVED = ("profile", "tol", "name", "extent", "dead_below", "dead_above", "interface_below", "interface_above") magnetic = True - def __init__(self, profile=None, tol=1e-3, name=None, **kw): + def __init__( + self, + profile_source=None, + profile=None, + tol=1e-3, + name=None, + thickness=None, + profile_params=None, + rhoM_start=None, + thetaM_start=None, + rhoM_end=None, + thetaM_end=None, + **kw, + ): + if profile_source is not None: + self.profile_source = profile_source + import numpy as np + import scipy + + ctx = {"np": np, "scipy": scipy} + output = {} + exec(profile_source, ctx, output) + # assume that the function definition is the only thing in the source + profile_name = list(output.keys())[0] + self.profile = output[profile_name] + elif profile is not None: + self.profile = profile + self.profile_source = inspect.getsource(profile) + profile_name = profile.__name__ + else: + raise TypeError("Need profile or profile_source") if not name: - name = profile.__name__ - if profile is None: - raise TypeError("Need profile") + name = profile_name + + self.name = name + self.tol = tol # strip magnetism keywords from list of keywords magkw = dict((a, kw.pop(a)) for a in set(self.RESERVED) & set(kw.keys())) BaseMagnetism.__init__(self, name=name, **magkw) - self.profile = profile - self.tol = tol + self.thickness = Parameter.default(thickness, name=name + " thickness") + self.rhoM_start = Parameter.default(rhoM_start, name=name + " rhoM_start") + self.thetaM_start = Parameter.default(thetaM_start, name=name + " thetaM_start") + self.rhoM_end = Parameter.default(rhoM_end, name=name + " rhoM_end") + self.thetaM_end = Parameter.default(thetaM_end, name=name + " thetaM_end") + + if profile_params is None: + # use kw arguments to set initial values for profile parameters + self.profile_params = _init_profile_params(name, self.profile, kw, self.RESERVED) + else: + # no need to filter out reserved keywords + self.profile_params = _init_profile_params(name, self.profile, profile_params, ()) + + self._set_ends() + + def set_anchor(self, stack: Stack, index): + stack, start = stack._lookup(index) + thickness_params = [] + for k in range(start, start + self.extent): + thickness_params.append(stack[k].thickness) + total_expression = sum(thickness_params) - self.dead_below - self.dead_above + self.thickness = total_expression + + def _set_ends(self): + rhoM_start = Calculation(self.name + " rhoM_start") + + def rhoM_start_func(): + P = self.profile(asarray([0.0]), **_get_values(self)) + return P[0][0] if isinstance(P, tuple) else P[0] + + rhoM_start.set_function(rhoM_start_func) + thetaM_start = Calculation(self.name + " thetaM_start") + + def thetaM_start_func(): + P = self.profile(asarray([0.0]), **_get_values(self)) + return P[1][0] if isinstance(P, tuple) else DEFAULT_THETA_M - _set_parameters(self, name, profile, kw, self.RESERVED) - rhoM_start = _MagnetismLimit(self, isend=False, isrhoM=True) - rhoM_end = _MagnetismLimit(self, isend=True, isrhoM=True) - thetaM_start = _MagnetismLimit(self, isend=False, isrhoM=False) - thetaM_end = _MagnetismLimit(self, isend=True, isrhoM=False) - self.start = Magnetism(rhoM=rhoM_start, thetaM=thetaM_start) - self.end = Magnetism(rhoM=rhoM_end, thetaM=thetaM_end) - self.anchor = None + thetaM_start.set_function(thetaM_start_func) + rhoM_end = Calculation(self.name + " rhoM_end") - def set_anchor(self, stack, index): - self.anchor = (stack, index) + def rhoM_end_func(): + P = self.profile(asarray([self.thickness.value]), **_get_values(self)) + return P[0][0] if isinstance(P, tuple) else P[0] + + rhoM_end.set_function(rhoM_end_func) + thetaM_end = Calculation(self.name + " thetaM_end") + + def thetaM_end_func(): + P = self.profile(asarray([self.thickness.value]), **_get_values(self)) + return P[1][0] if isinstance(P, tuple) else DEFAULT_THETA_M + + thetaM_end.set_function(thetaM_end_func) + self.rhoM_start.slot = rhoM_start + self.thetaM_start.slot = thetaM_start + self.rhoM_end.slot = rhoM_end + self.thetaM_end.slot = thetaM_end + + @property + def start(self): + return Magnetism(rhoM=self.rhoM_start, thetaM=self.thetaM_start) + + @property + def end(self): + return Magnetism(rhoM=self.rhoM_end, thetaM=self.thetaM_end) # TODO: is there a sane way of computing magnetism thickness in advance? def _calc_thickness(self): @@ -178,21 +333,9 @@ def _calc_thickness(self): def parameters(self): parameters = BaseMagnetism.parameters(self) - parameters.update(_get_parameters(self)) + parameters.update(self.profile_params) return parameters - def to_dict(self): - ret = BaseMagnetism.to_dict(self) - ret.update( - to_dict( - { - "profile": self.profile, - "parameters": _get_parameters(self), - "tol": self.tol, - } - ) - ) - def render(self, probe, slabs, thickness, anchor, sigma): Pw, Pz = slabs.microslabs(thickness) if len(Pw) == 0: @@ -214,94 +357,38 @@ def __repr__(self): return "FunctionalMagnetism(%s)" % self.name -def _set_parameters(self, name, profile, kw, reserved): +def _init_profile_params(name, profile, profile_params, reserved): # Query profile function for the list of arguments + output = {} vars = inspect.getfullargspec(profile)[0] # print "vars", vars if inspect.ismethod(profile): vars = vars[1:] # Chop self vars = vars[1:] # Chop z # print vars - unused = [k for k in kw.keys() if k not in vars] + unused = [k for k in profile_params.keys() if k not in vars] if len(unused) > 0: raise TypeError("Profile got unexpected keyword argument '%s'" % unused[0]) dups = [k for k in vars if k in reserved] if len(dups) > 0: raise TypeError("Profile has conflicting argument %r" % dups[0]) for k in vars: - kw.setdefault(k, 0) - for k, v in kw.items(): + output.setdefault(k, 0) + for k, v in profile_params.items(): try: pv = [Parameter.default(vi, name=f"{name} {k}[{i}]") for i, vi in enumerate(v)] except TypeError: pv = Parameter.default(v, name=f"{name} {k}") - setattr(self, k, pv) - self._parameters = vars - - -def _get_parameters(self): - return {k: getattr(self, k) for k in self._parameters} + output[k] = pv + return output def _get_values(self): vals = {} - for k in self._parameters: - v = getattr(self, k) + for k, v in self.profile_params.items(): + # v = getattr(self, k) if isinstance(v, list): vals[k] = asarray([vk.value for vk in v], "d") else: vals[k] = v.value return vals - - -class _LayerLimit(Calculation): - def __init__(self, flayer, isend=True, isrho=True): - self.description = f"{'rho' if isrho else 'irho'} Layer Limit, isend={isend}" - self.flayer = flayer - self.isend = isend - self.isrho = isrho - self.name = str(flayer) + self._tag - - @property - def _tag(self): - return (".rho_" if self.isrho else ".irho_") + ("end" if self.isend else "start") - - def parameters(self): - return [] - - def _function(self): - z = asarray([0.0, self.flayer.thickness.value]) - P = self.flayer.profile(asarray(z), **_get_values(self.flayer)) - index = 1 if self.isend else 0 - return real(P[index]) if self.isrho else imag(P[index]) - - def __repr__(self): - return repr(self.flayer) + self._tag - - -class _MagnetismLimit(Calculation): - def __init__(self, flayer, isend=True, isrhoM=True): - self.description = f"{'rhoM' if isrhoM else 'thetaM'} Magnetism Limit, isend={isend}" - self.flayer = flayer - self.isend = isend - self.isrhoM = isrhoM - self.name = str(flayer) + self._tag - - @property - def _tag(self): - return (".rhoM_" if self.isrhoM else ".thetaM_") + ("end" if self.isend else "start") - - def parameters(self): - return [] - - def _function(self): - zmax = self.flayer._calc_thickness() - z = asarray([0.0, zmax]) - P = self.flayer.profile(z, **_get_values(self.flayer)) - rhoM, thetaM = P if isinstance(P, tuple) else (P, DEFAULT_THETA_M) - rhoM, thetaM = [broadcast_to(v, z.shape) for v in (rhoM, thetaM)] - index = -1 if self.isend else 0 - return rhoM[index] if self.isrhoM else thetaM[index] - - def __repr__(self): - return repr(self.flayer) + self._tag