write() and _write_attr()

lisainstrument/hexagon.py

@@ -16,6 +16,8 @@ Authors:
 
 import numpy as np
 
+from h5py import File
+
 from . import noises
 
 
@@ -29,8 +31,8 @@ class Hexagon():
         locked_laser_asd (float): ASD added to locked lasers 2 & 3 [Hz/sqrt(Hz)]
     """
 
-    INDICES = [1, 2, 3]
-    BEATNOTES = [12, 23, 31]
+    INDICES = ['1', '2', '3']
+    BEATNOTES = ['12', '23', '31']
 
     def __init__(self,
                  size=100, dt=1 / (80e6/2**17/10/6/3),
@@ -39,6 +41,7 @@ class Hexagon():
                  offset_freqs='default',
                  central_freq=2.816E14):
 
+        self.simulated = False
         self.size = int(size)
         self.dt = float(dt)
         self.fs = 1.0 / self.dt
@@ -50,16 +53,16 @@ class Hexagon():
 
         self.central_freq = float(central_freq)
         if offset_freqs == 'default':
-            self.offset_freqs = {1: 0.0, 2: 15E6, 3: 7E6}
+            self.offset_freqs = {'1': 0.0, '2': 15E6, '3': 7E6}
         else:
             self.offset_freqs = offset_freqs
 
         # initialize single-laser-related attributes
         self.laser_noises = None
-        self.carrier_flucuations = None
         self.carrier_offsets = None
+        self.carrier_fluctuations = None
         # initialize beatnote-related attributes
-        self.carrier_beatnote = None
+        self.carrier_beatnotes = None
         self.carrier_beatnote_offsets = None
         self.carrier_beatnote_fluctuations = None
         # initialize signal-combination attributes
@@ -71,10 +74,11 @@ class Hexagon():
 
         Args:
         """
+        self.simulated = True
 
         self.laser_noises = np.empty((self.size, 3)) # (size, laser) [Hz]
 
-        # Laser 1 is free-running
+        # Laser 1 has its own stability
         self.laser_noises[:, 0] = noises.laser(self.fs, self.size, self.primary_laser_asd)
         # Laser 2 replicated laser 1 with added locking noise
         self.laser_noises[:, 1] = self.laser_noises[:, 0] \
@@ -84,7 +88,7 @@ class Hexagon():
             + noises.white(self.fs, self.size, self.locked_laser_asd)
 
         # Carrier beams
-        self.carrier_flucuations = self.laser_noises # (size, laser) [Hz]
+        self.carrier_fluctuations = self.laser_noises # (size, laser) [Hz]
         self.carrier_offsets = np.array(
             [[self.offset_freqs[index] for index in self.INDICES]]
         ) # (size, laser) [Hz]
@@ -92,15 +96,58 @@ class Hexagon():
         # Compute beatnotes
         # Convention is from paper: beatnote ij is beam j - beam i
         self.carrier_beatnote_offsets = np.stack([
-            self.carrier_offsets[:, int(str(ij)[1]) - 1] - self.carrier_offsets[:, int(str(ij)[0]) - 1]
+            self.carrier_offsets[:, int(ij[1]) - 1] - self.carrier_offsets[:, int(ij[0]) - 1]
             for ij in self.BEATNOTES
         ], axis=-1) # (size, beatnote) [Hz]
         self.carrier_beatnote_fluctuations = np.stack([
-            self.carrier_flucuations[:, int(str(ij)[1]) - 1] - self.carrier_flucuations[:, int(str(ij)[0]) - 1]
+            self.carrier_fluctuations[:, int(ij[1]) - 1] - self.carrier_fluctuations[:, int(ij[0]) - 1]
             for ij in self.BEATNOTES
         ], axis=-1) # (size, beatnote) [Hz]
-        self.carrier_beatnote = self.carrier_beatnote_offsets + self.carrier_beatnote_fluctuations
+        self.carrier_beatnotes = self.carrier_beatnote_offsets + self.carrier_beatnote_fluctuations
 
         # Three-signal combination
-        self.three_signal_combination = self.carrier_beatnote[:,0] \
-            + self.carrier_beatnote[:,1] + self.carrier_beatnote[:,2]
\ No newline at end of file
+        self.three_signal_combination = self.carrier_beatnotes[:,0] \
+            + self.carrier_beatnotes[:,1] + self.carrier_beatnotes[:,2]
+
+
+    def write(self, output='measurements.h5', mode='w'):
+        """Write simulation results.
+
+        Args:
+            output: path to measurement file
+            mode: measurement file opening mode
+        """
+
+        with File(output, mode) as hdf5:
+
+            if not self.simulated:
+                self.simulate()
+
+            self._write_attr(hdf5, 'laser_noises', indices=self.INDICES)
+            self._write_attr(hdf5, 'carrier_offsets', indices=self.INDICES)
+            self._write_attr(hdf5, 'carrier_fluctuations', indices=self.INDICES)
+            self._write_attr(hdf5, 'carrier_beatnotes', indices=self.BEATNOTES)
+            self._write_attr(hdf5, 'carrier_beatnote_offsets', indices=self.BEATNOTES)
+            self._write_attr(hdf5, 'carrier_beatnote_fluctuations', indices=self.BEATNOTES)
+            self._write_attr(hdf5, 'three_signal_combination')
+
+            
+    def _write_attr(self, hdf5, dataset, indices=None):
+        """Write a single object attribute on ``hdf5``.
+
+        Args:
+            hdf5 (:obj:`h5py.Group`): an HDF5 file, or a dataset
+            indices ([str]): index list
+        """
+
+        # Get a target attribute
+        attr = getattr(self, dataset)
+        # Write dataset
+        if indices is not None:
+            size = attr.shape[0]
+            dtype = np.dtype({'names': indices, 'formats': len(indices) * [np.float64]})
+            hdf5.create_dataset(dataset, (size,), dtype=dtype)
+            for i, index in enumerate(indices):
+                hdf5[dataset][index] = attr[:, i]
+        else:
+            hdf5.create_dataset(dataset, data=attr, dtype=np.float64)