from dataclasses import dataclass
import inspect
from typing import Dict, Optional, Callable, Union
import numpy as np
from numpy import asarray, broadcast_to, imag, real
from bumps.parameter import Calculation, Constant, Parameter, ValueProtocol
from refl1d import utils
from .layers import Layer, Stack
from .magnetism import DEFAULT_THETA_M, BaseMagnetism, Magnetism
from .material import SLD
[docs]
@dataclass
class FunctionalProfile(Layer):
"""
Generic profile function
Parameters:
*thickness* the thickness of the layer
*interface* the roughness of the surface [not implemented]
*profile* the profile function, suitably parameterized
*tol* is the tolerance for considering values equal
*magnetism* magnetic profile associated with the layer
*name* is the layer name
The profile function takes a depth vector *z* returns a density vector
*rho*. For absorbing profiles, return complex vector *rho + irho*1j*.
*z* is guaranteed to be increasing, with step size 2*z[0]. The
remaining profile arguments are added as fittable parameters to
the layer.
The profile parameters can be initialized using key=value arguments to
to FunctionalProfile. If the value is a sequence, then a list of parameters
will be created, one for each sequence element. These will be converted
to numpy vectors prior to calling the profile function. If no value is
supplied, the parameter will use the default name=value from the profile
function.
The attributes *start* and *end* are the SLD at at the start and end of
the profile respectively. The *rho* and *irho* attributes for these SLDs
can be used in parameter expressions for other layers. Note that each
one is implemented as a separate call to the profile function, using z=0
for start and z=thickness for end.
Within the profile function there is no mechanism for querying the larger
profile to determine the value of the *rho* at the layer boundaries. If
needed, this information will have to be communicated through shared
parameters. For example::
L1 = SLD('L1', rho=2.07)
L3 = SLD('L3', rho=4)
def linear(z, rhoL, rhoR):
rho = z * (rhoR-rhoL)/(z[-1]-z[0]) + rhoL
return rho
profile = FunctionalProfile(100, 0, profile=linear,
rhoL=L1.rho, rhoR=L3.rho)
sample = L1 | profile | L3
"""
thickness: Parameter
interface: Constant
profile: Callable = None
tol: float = 0
magnetism: Optional[BaseMagnetism] = None
name: str = ""
# Attributes required for serialize/deserialize
start: SLD = None
end: SLD = None
pars: Dict[str, Parameter] = None
# 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,
start=None,
end=None,
pars=None,
**kw,
):
# print(f"building FP with {thickness=} {interface=} {profile=} {name=} {start=} {end=} {pars=} {kw=}")
if not name:
name = profile.__name__
# TODO: maybe we already handle interface with the blend function?
if float(interface) != 0:
raise NotImplementedError("interface not yet supported")
self.name = name
self.thickness = Parameter.default(thickness, name=name + " thickness")
self.interface = Constant(float(interface), name=name + " interface")
self.profile = profile
self.tol = tol
self.magnetism = magnetism
self.start = start if start is not None else SLD(name + " start")
self.end = end if end is not None else SLD(name + " end")
self.pars = _parse_parameters(name, profile, kw) if pars is None else pars
# TODO: profile call is duplicated if asking for both rho and irho
# Fill calculation slots
self.start.rho.slot = Calculation(
description="profile rho at z=0",
function=lambda: real(self._eval([0.0])[0]),
)
self.start.irho.slot = Calculation(
description="profile irho at z=0",
function=lambda: imag(self._eval([0.0])[0]),
)
self.end.rho.slot = Calculation(
description="profile rho at z=thickness",
function=lambda: real(self._eval([float(self.thickness)])[0]),
)
self.end.irho.slot = Calculation(
description="profile irho at z=thickness",
function=lambda: imag(self._eval([float(self.thickness)])[0]),
)
# Allow dot access to members of the parameter dictionary. Existing attributes
# of the object take precedence.
def __setattr__(self, key, value):
# print(f"setting {key}")
if self.pars and key in self.pars and key not in self.__dict__:
if not isinstance(value, ValueProtocol):
raise TypeError("Can only assign parameter or expression to a parameter slot")
self.__dict__["pars"][key] = value
else:
super().__setattr__(key, value)
def __getattr__(self, key):
# print(f"getting {key}")
if self.pars and key in self.pars:
return self.pars[key]
raise AttributeError(f"{type(self)!r} has no attribute {key!r}")
def _eval(self, Pz):
args = {k: asarray([float(vi) for vi in v]) if isinstance(v, list) else float(v) for k, v in self.pars.items()}
# if self.profile is None:
# return np.full_like(Pz, np.nan, dtype=complex)
return asarray(self.profile(asarray(Pz), **args))
[docs]
def parameters(self):
# TODO: we are not including the calculated parameters
return {**self.pars, "thickness": self.thickness}
[docs]
def render(self, probe, slabs):
Pw, Pz = slabs.microslabs(self.thickness.value)
if len(Pw) == 0:
return
phi = self._eval(Pz)
if phi.shape != Pz.shape:
raise TypeError("profile function '%s' did not return array phi(z)" % self.profile.__name__)
Pw, phi = utils.merge_ends(Pw, phi, tol=self.tol)
# P = M*phi + S*(1-phi)
slabs.extend(rho=[real(phi)], irho=[imag(phi)], w=Pw)
# TODO: Is the following is sufficient for interfacial roughness?
# slabs.sigma[0] = float(self.interface)
[docs]
@dataclass
class FunctionalMagnetism(BaseMagnetism):
"""
Functional magnetism profile.
Parameters:
*profile* the profile function, suitably parameterized
*tol* is the tolerance for considering values equal
:class:`refl1d.sample.magnetism.BaseMagnetism` parameters
The profile function takes a depth vector *z* and returns a magnetism
vector *rhoM*. For magnetic twist, return a pair of vectors *(rhoM, thetaM)*.
Constants can be returned for *rhoM* or *thetaM*. If *thetaM* is not
provided it defaults to *thetaM=270*.
See :class:`FunctionalProfile` for a description of the the profile
function.
"""
profile: Callable
tol: float = 1e-3
# Attributes required for serialize/deserialize
start: Magnetism = None
end: Magnetism = None
thickness: Parameter = None
pars: Dict[str, Parameter] = None
BASE_PARS = ("extent", "dead_below", "dead_above", "interface_below", "interface_above")
magnetic = True
def __init__(
self,
profile=None,
tol=1e-3,
name=None,
start=None,
end=None,
thickness=None,
pars=None,
**kw,
):
# print(f"building FM with {profile=}, {tol=}, {name=}, {start=}, {end=}, {thickness=}, {pars=}, {kw=}")
if profile is None:
raise TypeError("Need profile")
if not name:
name = profile.__name__
# strip magnetism keywords from list of keywords
magkw = dict((a, kw.pop(a)) for a in set(self.BASE_PARS) & set(kw.keys()))
BaseMagnetism.__init__(self, name=name, **magkw)
self.profile = profile
self.tol = tol
self.start = start if start is not None else Magnetism(name=name + " start")
self.end = end if end is not None else Magnetism(name=name + " end")
self.thickness = Parameter.default(0, name=name + " thickness") if thickness is None else thickness
self.pars = _parse_parameters(name, profile, kw) if pars is None else pars
# TODO: profile call is duplicated if asking for both rhoM and thetaM
# Fill calculation slots
self.start.rhoM.slot = Calculation(
description="profile magnetic SLD at z=0",
function=lambda: self._eval([0.0])[0][0],
)
self.start.thetaM.slot = Calculation(
description="profile magnetic angle at z=0",
function=lambda: self._eval([0.0])[1][0],
)
self.end.rhoM.slot = Calculation(
description="profile magnetic SLD z=thickness",
function=lambda: self._eval([float(self.thickness)])[0][0],
)
self.end.thetaM.slot = Calculation(
description="profile magnetic angle at z=thickness",
function=lambda: self._eval([float(self.thickness)])[1][0],
)
# print(f"...FM {self.start=}, {self.end=}, {self.end.rhoM=}, {self.thickness=}")
[docs]
def set_anchor(self, stack: Stack, index: Union[str, int]):
"""
Rebuild thickness calculation whenever the sample stack changes. Called from
:func:`set_magnetism_anchors`.
"""
stack, start = stack._lookup(index)
total_thickness = sum(stack[k].thickness for k in range(start, start + self.extent))
self.thickness.equals(total_thickness - (self.dead_below + self.dead_above))
[docs]
def parameters(self):
# TODO: we are not including the calculated parameters
return {**BaseMagnetism.parameters(self), **self.pars}
# Allow dot access to members of the parameter dictionary. Existing attributes
# of the object take precedence.
def __setattr__(self, key, value):
# print(f"setting {key}")
if self.pars and key in self.pars and key not in self.__dict__:
if not isinstance(value, ValueProtocol):
raise TypeError("Can only assign parameter or expression to a parameter slot")
self.__dict__["pars"][key] = value
else:
super().__setattr__(key, value)
def __getattr__(self, key):
# print(f"getting {key}")
if self.pars and key in self.pars:
return self.pars[key]
raise AttributeError(f"{type(self)!r} has no attribute {key!r}")
def _eval(self, Pz):
args = {k: asarray([float(vi) for vi in v]) if isinstance(v, list) else float(v) for k, v in self.pars.items()}
Pz = asarray(Pz)
# if self.profile is None:
# return np.full_like(Pz, np.nan), np.full_like(Pz, np.nan)
P = self.profile(Pz, **args)
rhoM, thetaM = P if isinstance(P, tuple) else (P, DEFAULT_THETA_M)
try:
# rhoM or thetaM may be constant, lists or arrays (but not tuples!)
rhoM, thetaM = [broadcast_to(v, Pz.shape) for v in (rhoM, thetaM)]
except ValueError:
raise TypeError(f"Profile function for '{self.name}' returns incorrect shape")
return rhoM, thetaM
[docs]
def render(self, probe, slabs, thickness, anchor, sigma):
Pw, Pz = slabs.microslabs(thickness)
if len(Pw) == 0:
return
rhoM, thetaM = self._eval(Pz)
P = rhoM + thetaM * 0.001j # combine rhoM/thetaM so they can be merged
Pw, P = utils.merge_ends(Pw, P, tol=self.tol)
rhoM, thetaM = P.real, P.imag * 1000 # split out rhoM,thetaM again
slabs.add_magnetism(anchor=anchor, w=Pw, rhoM=rhoM, thetaM=thetaM, sigma=sigma)
def __repr__(self):
return "FunctionalMagnetism(%s)" % self.name
def _parse_parameters(name, profile, kw):
# Query profile function for the list of arguments
argspec = inspect.getfullargspec(profile)
vars = argspec.args
# 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]
if len(unused) > 0:
raise TypeError("Profile got unexpected keyword argument '%s'" % unused[0])
# Plug the default value from the function definition into any parameters
# not supplied by the caller.
defaults = argspec.defaults or ()
if len(defaults) > len(vars):
defaults = defaults[-len(vars) :]
if len(defaults) < len(vars):
defaults = [None] * (len(vars) - len(defaults)) + list(defaults)
for k, v in zip(vars, defaults):
kw.setdefault(k, v)
# Create the parameter dictionary.
pars = {}
for k, v in kw.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}")
pars[k] = pv
return pars
# TODO: Magnetism can't access the nuclear stack.
[docs]
def set_magnetism_anchors(stack):
"""
Find all magnetism objects in the stack and set their anchors.
This sets the total thickness, which is required to expected magnetism at the end
of the profile. This function needs to be called for samples that have FunctionalMagnetism
layers after the sample is built, and whenever the sample structure is changed.
"""
for k, layer in enumerate(stack):
if isinstance(layer.magnetism, FunctionalMagnetism):
layer.magnetism.set_anchor(stack, k)