diff --git a/attune/_arrangement.py b/attune/_arrangement.py index 25d38da..4ea3762 100644 --- a/attune/_arrangement.py +++ b/attune/_arrangement.py @@ -4,6 +4,7 @@ from typing import Dict, Union import numpy as np +import matplotlib.pyplot as plt from ._tune import Tune from ._discrete_tune import DiscreteTune @@ -89,6 +90,20 @@ def as_dict(self): out["tunes"] = {k: v.as_dict() for k, v in self._tunes.items()} return out + def plot(self): + nplots = len(self.tunes) + axes = [plt.subplot(nplots, 1, i + 1) for i in range(nplots)] + for i, (name, ti) in enumerate(self.tunes.items()): + axes[i].scatter(ti.independent, ti.dependent) + axes[i].plot(ti.independent, ti.dependent) + axes[i].grid() + axes[i].set_ylabel(f"{name} ({ti.dep_units})") + # can I assume all independent vars are the same? + if i + 1 != nplots: + plt.xticks(visible=False) + else: + axes[i].set_xlabel(f"wavelength ({ti.ind_units})") + @property def ind_max(self): """The maximum independant (input) value for this arrangement.""" diff --git a/attune/_instrument.py b/attune/_instrument.py index 1e58c16..86ff091 100644 --- a/attune/_instrument.py +++ b/attune/_instrument.py @@ -3,6 +3,7 @@ from datetime import datetime as _datetime from typing import Dict, Optional, Union +import matplotlib.pyplot as plt import json from ._arrangement import Arrangement @@ -164,6 +165,12 @@ def load(self): """The POSIX timestamp for when this instrument was created, if it was stored.""" return self._load + def plot(self, arrangement=None): + arrangements = self.arrangements if arrangement is None else arrangement + for ar in arrangements.values(): + plt.figure() + ar.plot() + def save(self, file): """Save the JSON representation into an open file.""" diff --git a/tests/instrument/test_plot.py b/tests/instrument/test_plot.py new file mode 100644 index 0000000..4e163f7 --- /dev/null +++ b/tests/instrument/test_plot.py @@ -0,0 +1,27 @@ +import attune +import numpy as np + + +def test_plot_arrangement(): + x = np.linspace(0, 2, 21) + a = attune.Tune(independent=x, dependent=x**0.5) + b = attune.Tune(independent=x, dependent=x**0.3) + arrangement = attune.Arrangement(name="test", tunes={"a": a, "b": b}) + arrangement.plot() + + +def test_plot_instrument(): + x = np.linspace(0, 2, 21) + a = attune.Tune(independent=x, dependent=x**0.5) + b = attune.Tune(independent=x, dependent=x**0.3) + arrangement = attune.Arrangement(name="test", tunes={"a": a, "b": b}) + instrument = attune.Instrument({"arr1": arrangement}, name="test") + instrument.plot() + + +if __name__ == "__main__": + import matplotlib.pyplot as plt + + test_plot_arrangement() + test_plot_instrument() + plt.show()