fix(serialization): use HDF5 for WFSC snippets

.github/workflows/pyfalco_test.yml

@@ -9,6 +9,7 @@ on:
   pull_request:
     branches:
       - master
+  workflow_dispatch: 
 
 jobs:
   build:
@@ -37,6 +38,7 @@ jobs:
         pip install pytest-html
         pip install numpy
         pip install scipy
+        pip install h5py
         pip install psutil
         pip install matplotlib
         pip install coveralls

falco/est.py

@@ -611,7 +611,7 @@ def pairwise_probing(mp, ev, jacStruct=np.array([])):
                         # If <2 probe pairs had good measurements, can't do pinv. Leave Eest as zero.
                         if NpairsGood < 2:
                             zerosCounter = zerosCounter + 1
-                            Epix = np.zeros((2, 1))
+                            Epix = np.zeros(2)
                             zerosCounter += 1
 
                         # Otherwise, use the 2+ good probe pair measurements for that pixel

falco/hdf5utils.py

@@ -0,0 +1,145 @@
+"""Helpers to persist FALCO objects in HDF5."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import types
+
+import h5py
+import numpy as np
+
+import falco
+
+
+_NODE_TYPE_ATTR = "node_type"
+
+
+def save_object_to_hdf5(obj, file_path: str | Path) -> None:
+    """Save a nested FALCO object tree to an HDF5 file.
+
+    Args:
+        obj: Object to serialize.
+        file_path: Destination HDF5 path.
+    """
+    with h5py.File(file_path, "w") as h5_file:
+        _write_node(h5_file, "root", obj)
+
+
+def load_object_from_hdf5(file_path: str | Path):
+    """Load an object tree produced by ``save_object_to_hdf5``.
+
+    Args:
+        file_path: Source HDF5 path.
+
+    Returns:
+        The reconstructed object.
+    """
+    with h5py.File(file_path, "r") as h5_file:
+        return _read_node(h5_file["root"])
+
+
+def _write_node(parent, name: str, value) -> None:
+    if isinstance(value, falco.config.Object):
+        group = parent.create_group(name)
+        group.attrs[_NODE_TYPE_ATTR] = "falco_object"
+        for key, sub_value in value.data.items():
+            _write_node(group, str(key), sub_value)
+        return
+
+    if isinstance(value, dict):
+        group = parent.create_group(name)
+        group.attrs[_NODE_TYPE_ATTR] = "dict"
+        for key, sub_value in value.items():
+            _write_node(group, str(key), sub_value)
+        return
+
+    if isinstance(value, (list, tuple)):
+        group = parent.create_group(name)
+        group.attrs[_NODE_TYPE_ATTR] = "tuple" if isinstance(value, tuple) else "list"
+        for index, sub_value in enumerate(value):
+            _write_node(group, f"item_{index:06d}", sub_value)
+        return
+
+    if isinstance(value, np.ndarray):
+        dataset = parent.create_dataset(name, data=value)
+        dataset.attrs[_NODE_TYPE_ATTR] = "ndarray"
+        return
+
+    if value is None:
+        group = parent.create_group(name)
+        group.attrs[_NODE_TYPE_ATTR] = "none"
+        return
+
+    if isinstance(value, str):
+        dataset = parent.create_dataset(
+            name, data=np.array(value, dtype=h5py.string_dtype("utf-8"))
+        )
+        dataset.attrs[_NODE_TYPE_ATTR] = "str"
+        return
+
+    if isinstance(value, bytes):
+        dataset = parent.create_dataset(name, data=np.frombuffer(value, dtype=np.uint8))
+        dataset.attrs[_NODE_TYPE_ATTR] = "bytes"
+        return
+
+    if isinstance(value, (np.generic, bool, int, float, complex)):
+        dataset = parent.create_dataset(name, data=np.asarray(value))
+        dataset.attrs[_NODE_TYPE_ATTR] = "scalar"
+        return
+
+    if isinstance(value, types.SimpleNamespace):
+        group = parent.create_group(name)
+        group.attrs[_NODE_TYPE_ATTR] = "simple_namespace"
+        for key, sub_value in vars(value).items():
+            _write_node(group, str(key), sub_value)
+        return
+
+    raise TypeError(f"Unsupported type for HDF5 serialization: {type(value)}")
+
+
+def _read_node(node):
+    if isinstance(node, h5py.Dataset):
+        node_type = node.attrs.get(_NODE_TYPE_ATTR, "dataset")
+        value = node[()]
+
+        if node_type == "bytes":
+            return bytes(np.asarray(value, dtype=np.uint8).tolist())
+
+        if isinstance(value, np.ndarray):
+            if node_type == "str":
+                return value.astype(str).item() if value.shape == () else value.astype(str)
+            return value
+
+        if isinstance(value, np.generic):
+            return value.item()
+
+        if isinstance(value, bytes):
+            return value.decode("utf-8") if node_type == "str" else value
+
+        return value
+
+    node_type = node.attrs.get(_NODE_TYPE_ATTR, "group")
+
+    if node_type == "none":
+        return None
+
+    if node_type == "falco_object":
+        obj = falco.config.Object()
+        for key in node.keys():
+            obj[key] = _read_node(node[key])
+        return obj
+
+    if node_type == "dict":
+        return {key: _read_node(node[key]) for key in node.keys()}
+
+    if node_type in ("list", "tuple"):
+        items = [_read_node(node[key]) for key in sorted(node.keys())]
+        return tuple(items) if node_type == "tuple" else items
+
+    if node_type == "simple_namespace":
+        obj = types.SimpleNamespace()
+        for key in node.keys():
+            setattr(obj, key, _read_node(node[key]))
+        return obj
+
+    return {key: _read_node(node[key]) for key in node.keys()}

falco/plot.py

@@ -308,7 +308,7 @@ def pairwise_probes(mp, ev, dDMVplus, ampSq2Dcube, iSubband):
         cmaps = Ncols*['gray']
         cmaps[0] = 'viridis'
         titles = ['DM Command (nm)', '+Probe Image', '-Probe Image',
-                  'Probe Intensity, $|\Delta p|^2$']
+              'Probe Intensity, $|\\Delta p|^2$']
 
         plusImageCube = ev.imageArray[:, :, 1::2, iSubband]
         minusImageCube = ev.imageArray[:, :, 2::2, iSubband]

falco/proper/prop_psd_errormap.py

@@ -72,12 +72,12 @@ def prop_psd_errormap(wf, amp, b, c, **kwargs):
     TPF : bool
         Indicates that the TPF 2D PSD shape is to be used:
 
-         .. math:: PSD\_2D(k) = \\frac{amp}{1 + (\\frac{k}{b})^c}
+         .. math:: PSD\\_2D(k) = \\frac{amp}{1 + (\\frac{k}{b})^c}
 
         where k and b are in cycles/meter. The default PSD shape is used if
         TPF is not specified:
 
-         .. math::  PSD\_2D(k) = \\frac{amp}{(1 + (\\frac{k}{b})^2)^{(c+1)/2}}
+         .. math::  PSD\\_2D(k) = \\frac{amp}{(1 + (\\frac{k}{b})^2)^{(c+1)/2}}
 
 
         This is the K-correlation form (see Church et al., Proc. of the SPIE,

falco/wfsc.py

@@ -7,6 +7,7 @@
 import matplotlib.pyplot as plt
 
 import falco
+from falco.hdf5utils import save_object_to_hdf5
 
 
 def loop(mp, out):
@@ -163,11 +164,10 @@ def loop(mp, out):
         else:
             print('Previous Measured NI:\t\t\t %.2e ' % (out.InormHist[Itr]))
 
-        # Save just the 'out' object to a pickle file
-        fnSnippet = os.path.join(mp.path.brief, (mp.runLabel + '_snippet.pkl'))
+        # Save just the 'out' object to an HDF5 file.
+        fnSnippet = os.path.join(mp.path.brief, (mp.runLabel + '_snippet.h5'))
         print('Saving data snippet to:\n\t%s  ...' % (fnSnippet), end='')
-        with open(fnSnippet, 'wb') as f:
-            pickle.dump(out, f)
+        save_object_to_hdf5(out, fnSnippet)
         print('done.', end='\n\n')
 
         # END OF ESTIMATION + CONTROL LOOP
@@ -189,11 +189,10 @@ def loop(mp, out):
         ev.Im = cvar.Im
         # plot_progress(mp, out, Itr, ImSimOffaxis, cvar.Im)
 
-    # Save just the 'out' object to a pickle file
-    fnSnippet = os.path.join(mp.path.brief, (mp.runLabel + '_snippet.pkl'))
+    # Save just the 'out' object to an HDF5 file.
+    fnSnippet = os.path.join(mp.path.brief, (mp.runLabel + '_snippet.h5'))
     print('\nSaving data snippet to:\n\t%s...' % (fnSnippet), end='')
-    with open(fnSnippet, 'wb') as f:
-        pickle.dump(out, f)
+    save_object_to_hdf5(out, fnSnippet)
     print('done.\n')
 
     # Save out the data from the workspace

requirements.txt

@@ -9,6 +9,7 @@ cycler==0.12.1
 deepmerge==2.0
 docutils==0.21.2
 fonttools==4.55.0
+h5py==3.12.1
 idna==3.10
 imagesize==1.4.1
 iniconfig==2.0.0

setup.py

@@ -96,6 +96,7 @@ def run_tests(self):
       'sphinx_rtd_theme', 
       'psutil', 
       'matplotlib',
+            'h5py',
       'numpy',
       'scipy',
 #       'proper',

tests/functional/test_wfsc_flc.py

@@ -5,6 +5,7 @@
 from math import isclose
 
 import falco
+import falco.hdf5utils as hdf5utils
 
 import config_wfsc_flc as CONFIG
 
@@ -23,6 +24,9 @@ def test_wfsc_flc():
     out = falco.setup.flesh_out_workspace(mp)
     falco.wfsc.loop(mp, out)
 
+    fn_hdf5 = os.path.join(mp.path.brief, f"{mp.runLabel}_snippet.h5")
+    out_from_hdf5 = hdf5utils.load_object_from_hdf5(fn_hdf5)
+
     print(out.IrawCorrHist[-1])
     print(out.IestScoreHist[-1])
     print(out.IincoCorrHist[-1])
@@ -51,6 +55,8 @@ def test_wfsc_flc():
     assert isclose(thput, 0.1486, abs_tol=1e-3)
 
     assert np.allclose(out.log10regHist, np.array([-2, -2, -2]), rtol=1e-2)
+    assert np.allclose(out_from_hdf5.log10regHist, out.log10regHist, rtol=1e-12)
+    assert isclose(out_from_hdf5.IrawCorrHist[-1], out.IrawCorrHist[-1], abs_tol=1e-14)
 
 
 if __name__ == '__main__':