Add instrument simulation

lisainstrument/containers.py

+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Container object for signals.
+
+Authors:
+    Jean-Baptiste Bayle <j2b.bayle@gmail.com>
+"""
+
+# TODO:
+#     * modulated beam, 4 args (carrier_offsets, carrier_fluctuations, usb_offsets, usb_fluctuations)
+#     * args can be scalars (repeated in dict), dict (check and convert to float), callable (call with mosa as arg)
+
+
+class Signal:
+    """Represent a signal expressed as frequency offsets and fluctuations."""
+
+    def __init__(self, offsets=0, fluctuations=0):
+        """Initialize a signal from frequency offsets and fluctuations.
+
+        Args:
+            offsets: frequency offsets [Hz]
+            fluctuations: frequency fluctuations [Hz]
+        """
+        if callable(offsets):
+            self.offsets = {mosa: offsets(mosa) for mosa in Instrument.MOSAS}
+        else:
+            self.offsets = Instrument.mosa_dict(offsets)
+
+        if callable(fluctuations):
+            self.fluctuations = {mosa: fluctuations(mosa) for mosa in Instrument.MOSAS}
+        else:
+            self.fluctuations = Instrument.mosa_dict(fluctuations)
+
+    def transformed(self, offsets=lambda x, mosa: x, fluctuations=lambda x, mosa: x):
+        """Return a new two-variable signal transforming offsets and fluctuations.
+
+        Args:
+            offsets: function of (offsets, mosa) returning new offsets [Hz]
+            fluctuations: function of (fluctuations, mosa) returning new fluctuations [Hz]
+
+        Returns:
+            A new `Signal` isntance where offsets and fluctuations have been transformed.
+        """
+        return self.__class__(
+            offsets={
+                mosa: offsets(self.offsets[mosa], mosa)
+                for mosa in Instrument.MOSAS
+            },
+            fluctuations={
+                mosa: fluctuations(self.fluctuations[mosa], mosa)
+                for mosa in Instrument.MOSAS
+            },
+        )
+
+    def reduce(self, function=lambda offsets, fluctuations: 0):
+        """Compute a new MOSA dictionary from offsets and fluctuations.
+
+        Args:
+            function: function of (offsets, fluctuations) returning new value
+        """
+        return {
+            mosa: function(self.offsets[mosa], self.fluctuations[mosa])
+            for mosa in Instrument.MOSAS
+        }
+
+    @property
+    def totalfreq(self):
+        """Return total frequencies, as the sum of offsets and fluctuations."""
+        return self.reduce(lambda offsets, fluctuations: offsets + fluctuations)
+
+
+class ModulatedBeam:
+    """Represent a modulated beam, with a carrier, an upper sideband, and optionally timer deviations."""
+
+    def __init__(self, carrier, usb, timer_deviations=None):
+        """Initialize a modulated beam from carrier and upper sideband signals.
+
+        Args:
+            carrier: carrier signal
+            usb: upper sideband signal
+            timer_deviations: timer deviations [s]
+        """
+        if not isinstance(carrier) or not isinstance(usb, Signal):
+            raise TypeError("carrier and upper sideband should be instances of `Signal`")
+
+        self.carrier = carrier
+        self.usb = usb
+        self.timer_deviations = Instrument.mosa_dict(timer_deviations)
+
+    def transformed(self,
+            carrier_offsets=lambda x, mosa: x, carrier_fluctuations=lambda x, mosa: x,
+            usb_offsets=lambda x, mosa: x, usb_fluctuations=lambda x, mosa: x,
+            timer_deviations=lambda x, mosa: x):
+        """Return a new modulated beam after applying transformations.
+
+        Args:
+            carrier_offsets: function of (offsets, mosa) returning new carrier offsets [Hz]
+            carrier_fluctuations: function of (fluctuations, mosa) returning new carrier fluctuations [Hz]
+            usb_offsets: function of (offsets, mosa) returning new upper sideband offsets [Hz]
+            usb_fluctuations: function of (fluctuations, mosa) returning new upper sideband fluctuations [Hz]
+            timer_deviations: function of (deviations, mosa) return new timer deviations [s]
+
+        Returns:
+            A new `ModulatedBeam` isntance where signals have been transformed.
+        """
+        return self.__class__(
+            carrier=self.carrier.transformed(carrier_offsets, carrier_fluctuations),
+            usb=self.usb.transformed(usb_offsets, usb_fluctuations),
+            timer_deviations={
+                mosa: timer_deviations(self.timer_deviations[mosa], mosa)
+                for mosa in Instrument.MOSAS
+            }
+        )