function to forward noises from sub-budget to an upper budget

gwinc/nb.py

@@ -98,6 +98,60 @@ def quadsum(data):
     return np.nansum(data, 0)
 
 
+def forward_noises(subbudgets):
+    """Extract noises and calibrations from a list of sub-budgets
+
+    Useful for forwarding a list of noises in a sub-budget into an upper level
+    budget. This then groups the noises in the upper level budget without
+    having to analyze the sub-budgets independently.
+
+    Parameters
+    ----------
+    subbudgets : list of Budgets
+      List of the sub-budgets whose noises will be forwarded.
+
+    Returns
+    -------
+    noises_frwd : list of Noise or (Noise, Calibration)
+      The list of noises and their calibrations to be forwarded
+
+    Example
+    -------
+    class Quantum(Budget):
+        noises = [
+            (Shot, Sensing),
+            RadiationPressure,
+        ]
+
+    class MainBudget(Budget):
+        noises = [
+            Thermal,
+        ]
+        noises += forward_noises([
+            Quantum,
+        ])
+
+    Defined like this the main budget will have shot noise, radiation pressure,
+    and thermal noise even though the two quantum noises are grouped in their
+    own sub-budget. If MainBudget had been defined with Quantum in the list of
+    noises with Thermal instead, the main budget would only have the quantum
+    and thermal noises, and the quantum noise would be further broken up
+    into shot noise and radiation pressure noises.
+    """
+    noises_frwd = []
+    for budget in subbudgets:
+        if not isinstance(budget, (tuple, list)):
+            budget = (budget,)
+        b = budget[0]
+        cals = tuple(budget[1:])
+        cals += tuple(b.calibrations)
+        noises_frwd.extend([
+            n + cals if isinstance(n, (tuple, list))
+            else (n,) + cals for n in b.noises
+        ])
+    return noises_frwd
+
+
 class BudgetItem:
     """GWINC BudgetItem class
 

test/budgets/H1/__init__.py

@@ -133,6 +133,13 @@ class SensingOpticalSpring(nb.Calibration):
         return 1 / np.abs(sensing_mA_m)**2
 
 
+class Quantum(nb.Budget):
+    noises = [
+        (Shot, Sensing),
+        RadiationPressure,
+    ]
+
+
 class DARMMeasured(nb.Noise):
     style = dict(label='H1 reference')
 
@@ -164,6 +171,7 @@ class H1(nb.Budget):
         (Shot, Sensing),
         RadiationPressure,
         Thermal,
+        noise.seismic.Seismic,
     ]
 
     references = [
@@ -177,14 +185,26 @@ class H1NoRefs(nb.Budget):
         (Shot, Sensing),
         RadiationPressure,
         Thermal,
+        noise.seismic.Seismic,
     ]
 
 
+class H1NoRefsForwardNoises(nb.Budget):
+    noises = [
+        noise.seismic.Seismic,
+    ]
+    noises += nb.forward_noises([
+        Quantum,
+        Thermal,
+    ])
+
+
 class H1dict(nb.Budget):
     noises = {
         'Shot': (Shot, Sensing),
         'RadiationPressure': RadiationPressure,
         'Thermal': ThermalDict,
+        'Seismic': noise.seismic.Seismic,
     }
 
     references = Struct(

test/budgets/test_budgets.py

@@ -164,6 +164,24 @@ def test_budget_dict_attributes(fpath_join, tpath_join):
     fig_dict1.savefig(tpath_join('budget_dict1.pdf'))
 
 
+@pytest.mark.logic
+def test_forward_noises(fpath_join, tpath_join, compare_noise):
+    B = load_budget(fpath_join('H1'), bname='H1NoRefs')
+    B_frwd = load_budget(fpath_join('H1'), bname='H1NoRefsForwardNoises')
+    tr = B.run()
+    tr_frwd = B_frwd.run()
+    compare_noise(tr.Thermal.SuspensionThermal, tr_frwd.SuspensionThermal)
+
+    fig = tr.plot()
+    fig.savefig(tpath_join('budget.pdf'))
+    fig_frwd = tr_frwd.plot()
+    fig_frwd.savefig(tpath_join('budget_frwd.pdf'))
+    fig = tr.Thermal.SuspensionThermal.plot()
+    fig.savefig(tpath_join('sus_thermal.pdf'))
+    fig = tr_frwd.SuspensionThermal.plot()
+    fig.savefig(tpath_join('sus_thermal_frwd.pdf'))
+
+
 @pytest.mark.logic
 @pytest.mark.fast
 def test_update_ifo_struct():