Source code for attune._intensity

"""Workup an instrument maximizing intensity of a single setable"""

import numpy as np
import WrightTools as wt

from ._instrument import Instrument
from ._arrangement import Arrangement
from ._setable import Setable
from ._tune import Tune
from ._transition import Transition
from ._plot import plot_intensity
from ._common import save


# --- processing methods --------------------------------------------------------------------------

__all__ = ["intensity"]


def _intensity(data, channel_name, tune_points, *, spline=True, **spline_kwargs):
    data.moment(axis=1, channel=channel_name, moment=1, resultant=data.axes[0].shape)
    offsets = data[f"{channel_name}_1_moment_1"].points

    if spline:
        spline = wt.kit.Spline(data.axes[0].points, offsets, **spline_kwargs)
        return spline(tune_points).clip(data.axes[1].min(), data.axes[1].max())
    if np.allclose(data.axes[0].points, tune_points):
        return offsets.clip(data.axes[1].min(), data.axes[1].max())
    elif np.allclose(data.axes[0].points, tune_points[::-1]):
        return offsets.clip(data.axes[1].min(), data.axes[1].max())[::-1]
    else:
        raise ValueError("Data points and instrument points do not match, and splining disabled")


[docs]def intensity( *, data, channel, arrangement, tune, instrument=None, level=False, gtol=0.01, ltol=0.1, autosave=True, save_directory=None, **spline_kwargs, ): """Workup a generic intensity plot for a single dependent. Parameters ---------- data : wt.data.Data should be in (setpoint, dependent) channel: wt.data.Channel or int or str channel to process arrangement: str name of the arrangement to modify in the instrument tune: str name of the tune to modify in the instrument instrument: attune.Instrument, optional instrument object to modify (Default None: make a new instrument) level: bool, optional toggle leveling data (Defalts to False) gtol: float, optional global tolerance for rejecting noise level relative to global maximum ltol: float, optional local tolerance for rejecting data relative to slice maximum autosave: bool, optional toggles saving of instrument file and images (Defaults to True) save_directory: Path-like where to save (Defaults to current working directory) **spline_kwargs: optional extra arguments to pass to spline creation (e.g. s=0, k=1 for linear interpolation) Returns ------- attune.Instrument New instrument object. """ metadata = { "channel": channel, "arrangement": arrangement, "tune": tune, "level": level, "gtol": gtol, "ltol": ltol, "spline_kwargs": spline_kwargs, } if not isinstance(channel, (int, str)): metadata["channel"] = channel.natural_name transition = Transition("intensity", instrument, metadata=metadata, data=data) data = data.copy() data.convert("nm") if instrument is not None: old_instrument = instrument.as_dict() setpoints = instrument[arrangement][tune].independent else: old_instrument = None setpoints = data.axes[0].points # TODO: units setpoints.sort() if isinstance(channel, (int, str)): channel = data.channels[wt.kit.get_index(data.channel_names, channel)] orig_channel = data.create_channel( f"{channel.natural_name}_orig", channel, units=channel.units ) # TODO: check if level does what we want if level: data.level(channel.natural_name, 0, -3) # TODO: gtol/ltol should maybe be moved to wt cutoff = channel.max() * gtol channel.clip(min=cutoff) max_axis = tuple(i for i, v in enumerate(data.axes[0].shape) if v > 1) cutoff = np.nanmax(channel[:], axis=1, keepdims=True) * ltol channel.clip(min=cutoff) offsets = _intensity(data, channel.natural_name, setpoints, **spline_kwargs) print(setpoints) print(offsets) try: raw_offsets = _intensity(data, channel.natural_name, setpoints, spline=False) except ValueError: raw_offsets = None units = data.axes[1].units if units == "None": units = None if instrument is not None: # Note, assumes units can be converted here... which I hope so if units is not None: offsets = wt.units.convert( offsets, units, old_instrument["arrangements"][arrangement]["tunes"][tune]["dep_units"], ) old_instrument["arrangements"][arrangement]["tunes"][tune]["independent"] = setpoints old_instrument["arrangements"][arrangement]["tunes"][tune]["dependent"] += offsets try: del old_instrument["transition"] except KeyError: pass new_instrument = Instrument(**old_instrument, transition=transition) else: arr = Arrangement(arrangement, {tune: Tune(setpoints, offsets, dep_units=units)}) new_instrument = Instrument( {arrangement: arr}, {tune: Setable(tune)}, transition=transition ) fig, _ = plot_intensity( data, channel.natural_name, arrangement, tune, new_instrument, instrument, raw_offsets ) if autosave: save(new_instrument, fig, "intensity", save_directory) return new_instrument