Merge pull request #48 from wright-group/setpoint

Initial setpoint implementation
wright-group · Jun 2, 2019 · 56a85f9 · 56a85f9
2 parents 37a194e + c649805
commit 56a85f9
Show file tree

Hide file tree

Showing 7 changed files with 274 additions and 2 deletions.
diff --git a/attune/workup/__init__.py b/attune/workup/__init__.py
@@ -1,3 +1,4 @@
 """Methods for processing OPA 800 tuning data."""
 from ._intensity import *
-from ._tune_test import *
+from ._setpoint import *
+from ._tune_test import *
diff --git a/attune/workup/_plot.py b/attune/workup/_plot.py
@@ -46,6 +46,51 @@ def plot_intensity(data, channel, dependent, curve, prior_curve=None, raw_offset
 
     return fig, gs
 
+def plot_setpoint(data, channel, dependent, curve, prior_curve=None, raw_offsets=None):
+    fig, gs = wt.artists.create_figure(
+        width="single", nrows=2, cols=[1, "cbar"], default_aspect=.5
+    )
+    ax = plt.subplot(gs[0, 0])
+    curve_plot_kwargs = {"lw": 5, "c": "k", "alpha": .5}
+    prior_curve_plot_kwargs = {"lw": 2, "c": "k"}
+    ax.plot(curve.setpoints[:], curve[dependent][:], **curve_plot_kwargs)
+    if prior_curve:
+        ax.plot(prior_curve.setpoints[:], prior_curve[dependent][:], **prior_curve_plot_kwargs)
+    ax.set_ylabel(dependent)
+    wt.artists.plot_gridlines()
+
+    ax = plt.subplot(gs[1, 0])
+    #data[channel][:] -= data.axes[0][:]
+    data[channel].signed = True
+    limits = -0.05*data[channel].mag(), 0.05 * data[channel].mag()
+    ax.pcolor(data, channel=channel, vmin=limits[0], vmax=limits[1])
+    xlim = ax.get_xlim()
+    if prior_curve:
+        ypoints = (
+            curve[dependent][:]
+            - prior_curve(curve.setpoints[:], curve.setpoints.units, full=False)[dependent]
+        )
+    else:
+        ypoints = curve[dependent][:]
+
+    if raw_offsets is not None:
+        ax.plot(curve.setpoints[:], raw_offsets, c="grey", lw=5, alpha=0.5)
+    ax.plot(curve.setpoints[:], ypoints, **curve_plot_kwargs)
+    ax.axhline(0, **prior_curve_plot_kwargs)
+    wt.artists.plot_gridlines()
+    ax.set_ylabel(r"$\mathsf{{\Delta {dependent}}}$".format(dependent=dependent))
+    ax.set_xlabel(f"Setpoint ({curve.setpoints.units})")
+    ax.set_xlim(xlim)
+
+    cax = plt.subplot(gs[1, 1])
+    ticks = np.linspace(*limits, 11)
+    wt.artists.plot_colorbar(
+        cax, clim=limits, ticks=ticks, label=channel, cmap="signed"
+    )
+
+    return fig, gs
+
+
 def plot_tune_test(data, channel, curve, prior_curve, raw_offsets=None):
     fig, gs = wt.artists.create_figure(default_aspect=0.5, cols=[1, "cbar"])
     # heatmap

diff --git a/attune/workup/_setpoint.py b/attune/workup/_setpoint.py
@@ -0,0 +1,113 @@
+"""Methods for processing OPA 800 tuning data."""
+
+import pathlib
+
+import numpy as np
+import WrightTools as wt
+
+from .. import Curve, Dependent, Setpoints
+from ._plot import plot_setpoint
+
+
+# --- processing methods --------------------------------------------------------------------------
+
+__all__ = ["setpoint"]
+
+
+def _setpoint(data, channel_name, tune_points, *, spline=True, **spline_kwargs):
+    offsets = []
+    chops = data.chop(1)
+    for c in chops.values():
+        xi = c.axes[0].points
+        yi = c[channel_name].points
+        xi, yi = wt.kit.remove_nans_1D(xi, yi)
+        if np.nanmin(yi) <= 0 <= np.nanmax(yi):
+            p = np.polynomial.Polynomial.fit(yi, xi, 2)
+            offsets.append(p(0))
+        else:
+            offsets.append(np.nan)
+
+    offsets = np.array(offsets)
+    if spline:
+        spline = wt.kit.Spline(data.axes[0].points, offsets, **spline_kwargs)
+        return spline(tune_points)
+    if np.allclose(data.axes[0].points, tune_points):
+        return offsets[::-1]
+    else:
+        raise ValueError("Data points and curve points do not match, and splining disabled")
+
+
+def setpoint(
+    data,
+    channel,
+    dependent,
+    curve=None,
+    *,
+    autosave=True,
+    save_directory=None,
+    **spline_kwargs,
+):
+    """Workup a generic intensity plot for a single dependent.
+
+    Parameters
+    ----------
+    data : wt.data.Data object
+        should be in (setpoint, dependent)
+
+    Returns
+    -------
+    curve
+        New curve object.
+    """
+    data = data.copy()
+    data.convert("wn")
+    if curve is not None:
+        old_curve = curve.copy()
+        old_curve.convert("wn")
+        setpoints = old_curve.setpoints
+    else:
+        old_curve = None   
+        setpoints = Setpoints(data.axes[0].points, data.axes[0].expression, data.axes[0].units)
+    # TODO: units
+
+    if isinstance(channel, (int, str)):
+        channel = data.channels[wt.kit.get_index(data.channel_names, channel)]
+
+    dims = [1] * data.ndim
+    dims[0] = setpoints[:].size # TODO: be more robust, don't assume 0 index
+    channel -= setpoints[:].reshape(dims)
+
+    offsets = _setpoint(data, channel.natural_name, setpoints[:], **spline_kwargs)
+    try:
+        raw_offsets = _setpoint(data, channel.natural_name, setpoints[:], spline=False)
+    except ValueError:
+        raw_offsets = None
+
+    units = data.axes[1].units
+    if units == "None":
+        units = None
+
+    new_curve = Curve(
+        setpoints, [Dependent(offsets, dependent, units, differential=True)], name="setpoint"
+    )
+
+    if curve is not None:
+        curve = old_curve + new_curve
+    else:
+        curve = new_curve
+
+    # Why did we have to map setpoints?
+    curve.map_setpoints(setpoints[:])
+
+    fig, _ = plot_setpoint(data, channel.natural_name, dependent, curve, old_curve, raw_offsets)
+
+    if autosave:
+        if save_directory is None:
+            # TODO: Formal decision on whether this should be cwd or data/curve location
+            save_directory = "."
+        save_directory = pathlib.Path(save_directory)
+        curve.save(save_directory=save_directory, full=True)
+        # Should we timestamp the image?
+        p = save_directory / "setpoint.png"
+        wt.artists.savefig(p, fig=fig)
+    return curve
diff --git a/setup.py b/setup.py
@@ -30,7 +30,7 @@ def read(fname):
     python_requires=">=3.5",
     setup_requires=["pytest-runner"],
     tests_require=["pytest", "pytest-cov"],
-    install_requires=["WrightTools>=3.0", "numpy", "scipy", "matplotlib", "tidy_headers"],
+    install_requires=["WrightTools>=3.2.2", "numpy", "scipy", "matplotlib", "tidy_headers"],
     extras_require={"dev": ["black", "pre-commit"], "docs": ["sphinx-gallery>0.3.0", "sphinx", "sphinx-rtd-theme"]},
     version=version,
     description="Tools for tuning optical parametric amplifiers and multidimensional spectrometers.",

diff --git a/tests/workup/setpoint/2018-11-30/aixcb.py b/tests/workup/setpoint/2018-11-30/aixcb.py
@@ -0,0 +1,37 @@
+import attune
+import pathlib
+import pytest
+import WrightTools as wt
+
+import matplotlib.pyplot as plt
+
+
+__here__ = pathlib.Path(__file__).parent
+
+
+def test():
+    data = wt.open(__here__ / "data.wt5")
+    data.print_tree()
+
+    data.convert("wn", convert_variables=True)
+    data.transform("w1=wm", "w1_Crystal_2_points", "wa-w1")
+    data.level(0,2,5)
+    data.array_signal.clip(min=0)
+    data.transform("w1=wm", "w1_Crystal_2_points", "wa")
+    data.moment("wa", moment=1, resultant=wt.kit.joint_shape(data.w1, data.w1_Crystal_2))
+    data.transform("w1=wm", "w1_Crystal_2_points")
+    data.channels[-1].clip(min=data.w1.min()-1000, max=data.w1.max()+1000)
+    data.channels[-1].null = data.wa.min()
+
+    old = attune.TopasCurve.read(
+        [__here__ / "old.crv"],
+        interaction_string="NON-NON-NON-Sig",
+    )
+
+    new = attune.workup.setpoint(data, -1, "2", autosave=False, curve=old)
+
+    print(new)
+
+
+if __name__ == "__main__":
+    test()
diff --git a/tests/workup/setpoint/2018-11-30/data.wt5 b/tests/workup/setpoint/2018-11-30/data.wt5
diff --git a/tests/workup/setpoint/2018-11-30/old.crv b/tests/workup/setpoint/2018-11-30/old.crv
@@ -0,0 +1,76 @@
+600						
+OPA/NOPA						
+0						
+0						
+-1						
+0						
+0						
+4						
+0	1	2	3			
+2						
+NON-NON-NON-Sig						
+1						
+
+1						
+795						
+4						
+0	0	0	0			
+25
+795.000000	1140.000000	4	-9.324437	2.165817	7.259192	56.780087	
+795.000000	1160.000000	4	-9.160146	2.187580	7.040502	56.874502	
+795.000000	1180.000000	4	-8.986959	2.208514	6.840363	56.963150	
+795.000000	1200.000000	4	-8.804875	2.228620	6.647657	57.046032	
+795.000000	1220.000000	4	-8.613894	2.247896	6.454240	57.123147	
+795.000000	1240.000000	4	-8.414017	2.266344	6.254239	57.194495	
+795.000000	1260.000000	4	-8.205244	2.283963	6.042106	57.260077	
+795.000000	1280.000000	4	-7.987574	2.300754	5.814651	57.319892	
+795.000000	1300.000000	4	-7.761008	2.316715	5.572952	57.373940	
+795.000000	1320.000000	4	-7.525545	2.331848	5.320444	57.422221	
+795.000000	1340.000000	4	-7.281186	2.346152	5.061374	57.464736	
+795.000000	1360.000000	4	-7.027931	2.359628	4.799532	57.501484	
+795.000000	1380.000000	4	-6.765779	2.372275	4.535686	57.532465	
+795.000000	1400.000000	4	-6.494730	2.384092	4.267147	57.557679	
+795.000000	1420.000000	4	-6.214785	2.395082	3.990444	57.577127	
+795.000000	1440.000000	4	-5.925944	2.405242	3.701662	57.590808	
+795.000000	1460.000000	4	-5.628206	2.414574	3.396541	57.598722	
+795.000000	1480.000000	4	-5.321572	2.423076	3.074929	57.600870	
+795.000000	1500.000000	4	-5.006041	2.430751	2.747495	57.597250	
+795.000000	1520.000000	4	-4.681614	2.437596	2.426308	57.587864	
+795.000000	1540.000000	4	-4.348291	2.443613	2.119391	57.572712	
+795.000000	1560.000000	4	-4.006070	2.448801	1.830076	57.551792	
+795.000000	1580.000000	4	-3.654954	2.453160	1.562761	57.525106	
+795.000000	1600.000000	4	-3.294941	2.456690	1.313887	57.492653	
+795.000000	1620.000000	4	-2.926032	2.459392	1.079490	57.454434	
+NON-NON-NON-Idl						
+1						
+
+0						
+795						
+4						
+0	0	0	0			
+25
+795.000000	1561.090909	4	-2.926032	2.459392	1.079490	57.454434	
+795.000000	1580.124224	4	-3.294941	2.456690	1.313887	57.492653	
+795.000000	1600.127389	4	-3.654954	2.453160	1.562761	57.525106	
+795.000000	1621.176471	4	-4.006070	2.448801	1.830076	57.551792	
+795.000000	1643.355705	4	-4.348291	2.443613	2.119391	57.572712	
+795.000000	1666.758621	4	-4.681614	2.437596	2.426308	57.587864	
+795.000000	1691.489362	4	-5.006041	2.430751	2.747495	57.597250	
+795.000000	1717.664234	4	-5.321572	2.423076	3.074929	57.600870	
+795.000000	1745.413534	4	-5.628206	2.414574	3.396541	57.598722	
+795.000000	1774.883721	4	-5.925944	2.405242	3.701662	57.590808	
+795.000000	1806.240000	4	-6.214785	2.395082	3.990444	57.577127	
+795.000000	1839.669421	4	-6.494730	2.384092	4.267147	57.557679	
+795.000000	1875.384615	4	-6.765779	2.372275	4.535686	57.532465	
+795.000000	1913.628319	4	-7.027931	2.359628	4.799532	57.501484	
+795.000000	1954.678899	4	-7.281186	2.346152	5.061374	57.464736	
+795.000000	1998.857143	4	-7.525545	2.331848	5.320444	57.422221	
+795.000000	2046.534653	4	-7.761008	2.316715	5.572952	57.373940	
+795.000000	2098.144330	4	-7.987574	2.300754	5.814651	57.319892	
+795.000000	2154.193548	4	-8.205244	2.283963	6.042106	57.260077	
+795.000000	2215.280899	4	-8.414017	2.266344	6.254239	57.194495	
+795.000000	2282.117647	4	-8.613894	2.247896	6.454240	57.123147	
+795.000000	2355.555556	4	-8.804875	2.228620	6.647657	57.046032	
+795.000000	2436.623377	4	-8.986959	2.208514	6.840363	56.963150	
+795.000000	2526.575342	4	-9.160146	2.187580	7.040502	56.874502	
+795.000000	2626.956522	4	-9.324437	2.165817	7.259192	56.780087