preselect amplifiers based on restrictions and bands

make sure that selected amplifiers (single or multiband) have a band
that encompasses design_band.

Signed-off-by: EstherLerouzic <[email protected]>
Change-Id: I8b66755efbe8413f32328b9e02099ffdedd4b7ed

gnpy/core/network.py

@@ -440,8 +440,8 @@ def compute_gain_power_target(node: elements.Edfa, prev_node, next_node, power_m
 
 
 def set_one_amplifier(node: elements.Edfa, prev_node, next_node, power_mode: bool, prev_voa: float, prev_dp: float,
-                      pref_ch_db: float, pref_total_db: float, network: DiGraph, equipment: dict, verbose: bool) \
-        -> Tuple[float, float]:
+                      pref_ch_db: float, pref_total_db: float, network: DiGraph, restrictions: List[str],
+                      equipment: dict, verbose: bool) -> Tuple[float, float]:
     """Set the EDFA amplifier configuration based on power targets:
 
     This function adjusts the amplifier settings according to the specified parameters and
@@ -458,6 +458,7 @@ def set_one_amplifier(node: elements.Edfa, prev_node, next_node, power_mode: boo
         pref_ch_db (float): reference per channel power in dB.
         pref_total_db (float): reference total power in dB.
         network (DiGraph): The network graph.
+        restrictions: (List[str]): The list of amplifiers authorized for this configuration.
         equipment (dict): Equipment library.
         verbose (bool): Flag for verbose logging.
 
@@ -475,16 +476,6 @@ def set_one_amplifier(node: elements.Edfa, prev_node, next_node, power_mode: boo
         raman_allowed = False
 
     if node.params.type_variety == '':
-        if node.variety_list and isinstance(node.variety_list, list):
-            restrictions = node.variety_list
-        elif isinstance(prev_node, elements.Roadm) and prev_node.restrictions['booster_variety_list']:
-            # implementation of restrictions on roadm boosters
-            restrictions = prev_node.restrictions['booster_variety_list']
-        elif isinstance(next_node, elements.Roadm) and next_node.restrictions['preamp_variety_list']:
-            # implementation of restrictions on roadm preamp
-            restrictions = next_node.restrictions['preamp_variety_list']
-        else:
-            restrictions = None
         edfa_eqpt = {n: a for n, a in equipment['Edfa'].items() if a.type_def != 'multi_band'}
         edfa_variety, power_reduction = \
             select_edfa(raman_allowed, gain_target, power_target, edfa_eqpt,
@@ -493,6 +484,7 @@ def set_one_amplifier(node: elements.Edfa, prev_node, next_node, power_mode: boo
                         restrictions=restrictions, verbose=verbose)
         extra_params = equipment['Edfa'][edfa_variety]
         node.params.update_params(extra_params.__dict__)
+        node.type_variety = node.params.type_variety
         dp += power_reduction
         gain_target += power_reduction
     else:
@@ -544,6 +536,60 @@ def set_one_amplifier(node: elements.Edfa, prev_node, next_node, power_mode: boo
     return dp, voa
 
 
+def get_node_restrictions(node: Union[elements.Edfa, elements.Multiband_amplifier], prev_node,
+                          next_node, equipment: dict, _design_bands: dict) -> List:
+    """Returns a list of eligible amplifiers that comply with restrictions and design bands.
+
+    If the node is a multiband amplifier, only multiband amplifiers will be considered.
+
+    Args:
+        node (Union[elements.Edfa, elements.Multiband_amplifier]): The current amplifier node.
+        prev_node: The previous node in the network.
+        next_node: The next node in the network.
+        equipment (Dict): A dictionary containing equipment specifications.
+        _design_bands (Dict): A dictionary of design bands with frequency limits.
+
+    Returns:
+        List[str]: A list of eligible amplifier types that meet the specified restrictions.
+    """
+    if node.params.type_variety != '' and node.params.type_variety:
+        # type_variety takes precedence over any other restrictions
+        return [node.params.type_variety]
+    restrictions = []
+    if node.variety_list and isinstance(node.variety_list, list):
+        restrictions = node.variety_list
+    elif isinstance(prev_node, elements.Roadm) and prev_node.restrictions['booster_variety_list']:
+        # implementation of restrictions on roadm boosters
+        restrictions = prev_node.restrictions['booster_variety_list']
+    elif isinstance(next_node, elements.Roadm) and next_node.restrictions['preamp_variety_list']:
+        # implementation of restrictions on roadm preamp
+        restrictions = next_node.restrictions['preamp_variety_list']
+    if isinstance(node, elements.Multiband_amplifier):
+        # Only keep multiband amps that are eligible for all the bands
+        # use the subset of EDFA library that are multiband, fits the design band and are either imposed
+        # in restriction list or allowed for design
+        multiband_eqpt = [n for n, a in equipment['Edfa'].items()
+                          if a.type_def == 'multi_band'
+                          and (n in restrictions or (not restrictions and a.allowed_for_design))]
+        # collect the individual amps part of the multiband amps that match the bands
+        edfa_eqpt = [t for m in multiband_eqpt
+                     for t in equipment['Edfa'][m].multi_band
+                     for band in _design_bands.values()
+                     if equipment['Edfa'][t].f_min <= band['f_min']
+                     and equipment['Edfa'][t].f_max >= band['f_max']]
+        # then filter all multi band amps whose amps group belong to the previous list
+        multiband_eqpt = [m for m in multiband_eqpt if all(t in edfa_eqpt for t in equipment['Edfa'][m].multi_band)]
+        # and returns the list of type_variety of multiband amps built with this single band amps
+        return multiband_eqpt
+    if isinstance(node, elements.Edfa):
+        band = next(b for b in _design_bands.values())
+        # preselect amps which are either part of restrictions or allowed for design, and compliant to the band.
+        edfa_eqpt = [n for n, a in equipment['Edfa'].items()
+                     if (a.type_def != 'multi_band' and a.f_min <= band['f_min'] and a.f_max >= band['f_max'])
+                     and (n in restrictions or (not restrictions and a.allowed_for_design))]
+        return edfa_eqpt
+
+
 def set_egress_amplifier(network: DiGraph, this_node: Union[elements.Roadm, elements.Transceiver], equipment: dict,
                          pref_ch_db: float, pref_total_db: float, verbose: bool):
     """This node can be a transceiver or a ROADM (same function called in both cases).
@@ -600,20 +646,31 @@ def set_egress_amplifier(network: DiGraph, this_node: Union[elements.Roadm, elem
             # go through all nodes in the OMS (loop until next Roadm instance)
             if isinstance(node, elements.Edfa):
                 band_name, _ = next((n, b) for n, b in _design_bands.items())
+                restrictions = get_node_restrictions(node, prev_node, next_node, equipment, _design_bands)
                 dp[band_name], voa[band_name] = set_one_amplifier(node, prev_node, next_node, power_mode,
                                                                   prev_voa[band_name], prev_dp[band_name],
                                                                   pref_ch_db, pref_total_db,
-                                                                  network, equipment, verbose)
+                                                                  network, restrictions, equipment, verbose)
             elif isinstance(node, elements.RamanFiber):
                 # this is to record the expected gain in Raman fiber in its .estimated_gain attribute.
                 band_name, _ = next((n, b) for n, b in _design_bands.items())
                 _ = span_loss(network, node, equipment, input_power=pref_ch_db + dp[band_name])
             elif isinstance(node, elements.Multiband_amplifier):
+                if len(node.amplifiers) == 0:
+                    # creates one amp per design band.
+                    for band_name, band in _design_bands.items():
+                        node.amplifiers[band_name] = elements.Edfa(params=EdfaParams.default_values, uid=node.uid)
+                # only select amplifiers which match the design bands
+                restrictions_multi = get_node_restrictions(node, prev_node, next_node, equipment, _design_bands)
+                restrictions_edfa = [e for m in restrictions_multi for e in equipment['Edfa'][m].multi_band]
                 for band_name, amp in node.amplifiers.items():
+                    _restrictions = [n for n in restrictions_edfa
+                                     if equipment['Edfa'][n].f_min <= _design_bands[band_name]['f_min']
+                                     and equipment['Edfa'][n].f_max >= _design_bands[band_name]['f_max']]
                     dp[band_name], voa[band_name] = \
                         set_one_amplifier(amp, prev_node, next_node, power_mode,
                                           prev_voa[band_name], prev_dp[band_name],
-                                          pref_ch_db, pref_total_db, network, equipment, verbose)
+                                          pref_ch_db, pref_total_db, network, _restrictions, equipment, verbose)
             prev_dp.update(**dp)
             prev_voa.update(**voa)
             prev_node = node

tests/test_network_functions.py

@@ -7,10 +7,11 @@
 from pathlib import Path
 import pytest
 from gnpy.core.exceptions import NetworkTopologyError
-from gnpy.core.network import span_loss, build_network, select_edfa
+from gnpy.core.network import span_loss, build_network, select_edfa, get_node_restrictions
 from gnpy.tools.json_io import load_equipment, load_network, network_from_json
-from gnpy.core.utils import lin2db, automatic_nch
-from gnpy.core.elements import Fiber, Edfa
+from gnpy.core.utils import lin2db, automatic_nch, merge_amplifier_restrictions
+from gnpy.core.elements import Fiber, Edfa, Roadm, Multiband_amplifier
+from gnpy.core.parameters import EdfaParams, MultiBandParams
 
 
 TEST_DIR = Path(__file__).parent
@@ -421,7 +422,10 @@ def network_base(case, site_type, length=50.0, amplifier_type='Multiband_amplifi
      [{'f_min': 191.3e12, 'f_max': 196.1e12}], [{'f_min': 191.3e12, 'f_max': 196.1e12}]),
     ('design', 'Fused', 'Multiband_amplifier',
      [{'f_min': 191.3e12, 'f_max': 196.1e12}],
-     [{'f_min': 186.55e12, 'f_max': 190.05e12}, {'f_min': 191.25e12, 'f_max': 196.15e12}])])
+     [{'f_min': 186.55e12, 'f_max': 190.05e12}, {'f_min': 191.25e12, 'f_max': 196.15e12}]),
+    ('no_design', 'Fused', 'Multiband_amplifier',
+     [{'f_min': 191.3e12, 'f_max': 196.1e12}],
+     [{'f_min': 187.0e12, 'f_max': 190.0e12}, {'f_min': 191.3e12, 'f_max': 196.0e12}])])
 def test_design_band(case, site_type, amplifier_type, expected_design_bands, expected_per_degree_design_bands):
     """Check design_band is the one defined:
     - in SI if nothing is defined,
@@ -461,3 +465,80 @@ def test_select_edfa(caplog, raman_allowed, gain_target, power_target, target_ex
     assert selection == expected_selection
     if warning:
         assert warning in caplog.text
+
+
+@pytest.mark.parametrize('cls, defaultparams, variety_list, booster_list, band, expected_restrictions', [
+    (Edfa, EdfaParams, [], [],
+     {'LBAND': {'f_min': 187.0e12, 'f_max': 190.0e12}},
+     ['std_medium_gain_L', 'std_low_gain_L_ter', 'std_low_gain_L']),
+    (Edfa, EdfaParams, [], [],
+     {'CBAND': {'f_min': 191.3e12, 'f_max': 196.0e12}},
+     ['CienaDB_medium_gain', 'std_medium_gain', 'std_low_gain', 'std_low_gain_bis', 'test', 'test_fixed_gain']),
+    (Edfa, EdfaParams, ['std_medium_gain', 'std_high_gain'], [],
+     {'CBAND': {'f_min': 191.3e12, 'f_max': 196.0e12}},
+     ['std_medium_gain']),   # name in variety list does not exist in library
+    (Edfa, EdfaParams, ['std_medium_gain', 'std_high_gain'], [],
+     {'LBAND': {'f_min': 187.0e12, 'f_max': 190.0e12}},
+     []),   # restrictions inconsistency with bands
+    (Edfa, EdfaParams, ['std_medium_gain', 'std_high_gain'], ['std_booster'],
+     {'CBAND': {'f_min': 191.3e12, 'f_max': 196.0e12}},
+     ['std_medium_gain']),  # variety list takes precedence over booster constraint
+    (Edfa, EdfaParams, [], ['std_booster'],
+     {'CBAND': {'f_min': 191.3e12, 'f_max': 196.0e12}},
+     ['std_booster']),
+    (Multiband_amplifier, MultiBandParams, [], [],
+     {'CBAND': {'f_min': 191.3e12, 'f_max': 196.0e12}, 'LBAND': {'f_min': 187.0e12, 'f_max': 190.0e12}},
+     ['std_medium_gain_multiband', 'std_low_gain_multiband_bis']),
+    (Multiband_amplifier, MultiBandParams, [], ['std_booster_multiband', 'std_booster'],
+     {'CBAND': {'f_min': 191.3e12, 'f_max': 196.0e12}, 'LBAND': {'f_min': 187.0e12, 'f_max': 190.0e12}},
+     ['std_booster_multiband'])
+])
+def test_get_node_restrictions(cls, defaultparams, variety_list, booster_list, band, expected_restrictions):
+    """Check that all combinations of restrictions are correctly captured
+    """
+    equipment = load_equipment(EQPT_MULTBAND_FILENAME)
+    edfa_config = {"uid": "Edfa1"}
+    if cls == Multiband_amplifier:
+        edfa_config['amplifiers'] = {}
+    edfa_config['params'] = defaultparams.default_values
+    edfa_config['variety_list'] = variety_list
+    node = cls(**edfa_config)
+    roadm_config = {
+        "uid": "roadm Brest_KLA",
+        "params": {
+            "per_degree_pch_out_db": {},
+            "target_pch_out_dbm": -18,
+            "add_drop_osnr": 38,
+            "pmd": 0,
+            "pdl": 0,
+            "restrictions": {
+                "preamp_variety_list": [],
+                "booster_variety_list": booster_list
+            },
+            "roadm-path-impairments": []
+        },
+        "metadata": {
+            "location": {
+                "city": "Brest_KLA",
+                "region": "RLD",
+                "latitude": 4.0,
+                "longitude": 0.0
+            }
+        }
+    }
+    prev_node = Roadm(**roadm_config)
+    fiber_config = {
+        "uid": "fiber (SITE1 → ILA1)",
+        "type_variety": "SSMF",
+        "params": {
+            "length": 100.0,
+            "loss_coef": 0.2,
+            "length_units": "km"
+        }
+    }
+    extra_params = equipment['Fiber']['SSMF'].__dict__
+
+    fiber_config['params'] = merge_amplifier_restrictions(fiber_config['params'], extra_params)
+    next_node = Fiber(**fiber_config)
+    restrictions = get_node_restrictions(node, prev_node, next_node, equipment, band)
+    assert restrictions == expected_restrictions