Environments

EnvParams

Dataclass to hold environment parameters. Parameters are immutable.

Parameters:

Name	Type	Description	Default
max_requests	Scalar	Maximum number of requests in an episode	required
incremental_loading	Scalar	Incremental increase in traffic load (non-expiring requests)	required
end_first_blocking	Scalar	End episode on first blocking event	required
continuous_operation	Scalar	If True, do not reset the environment at the end of an episode	required
edges	Array	Two column array defining source-dest node-pair edges of the graph	required
slot_size	Scalar	Spectral width of frequency slot in GHz	required
consider_modulation_format	Scalar	If True, consider modulation format to determine required slots	required
link_length_array	Array	Array of link lengths	required
aggregate_slots	Scalar	Number of slots to aggregate into a single action (First-Fit with aggregation)	required
guardband	Scalar	Guard band in slots	required
directed_graph	bool	Whether graph is directed (one fibre per link per transmission direction)	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvParams:
    """Dataclass to hold environment parameters. Parameters are immutable.

    Args:
        max_requests (chex.Scalar): Maximum number of requests in an episode
        incremental_loading (chex.Scalar): Incremental increase in traffic load (non-expiring requests)
        end_first_blocking (chex.Scalar): End episode on first blocking event
        continuous_operation (chex.Scalar): If True, do not reset the environment at the end of an episode
        edges (chex.Array): Two column array defining source-dest node-pair edges of the graph
        slot_size (chex.Scalar): Spectral width of frequency slot in GHz
        consider_modulation_format (chex.Scalar): If True, consider modulation format to determine required slots
        link_length_array (chex.Array): Array of link lengths
        aggregate_slots (chex.Scalar): Number of slots to aggregate into a single action (First-Fit with aggregation)
        guardband (chex.Scalar): Guard band in slots
        directed_graph (bool): Whether graph is directed (one fibre per link per transmission direction)
    """
    max_requests: chex.Scalar = struct.field(pytree_node=False)
    incremental_loading: chex.Scalar = struct.field(pytree_node=False)
    end_first_blocking: chex.Scalar = struct.field(pytree_node=False)
    continuous_operation: chex.Scalar = struct.field(pytree_node=False)
    edges: chex.Array = struct.field(pytree_node=False)
    slot_size: chex.Scalar = struct.field(pytree_node=False)
    consider_modulation_format: chex.Scalar = struct.field(pytree_node=False)
    link_length_array: chex.Array = struct.field(pytree_node=False)
    aggregate_slots: chex.Scalar = struct.field(pytree_node=False)
    guardband: chex.Scalar = struct.field(pytree_node=False)
    directed_graph: bool = struct.field(pytree_node=False)
    maximise_throughput: bool = struct.field(pytree_node=False)
    reward_type: str = struct.field(pytree_node=False)
    values_bw: chex.Array = struct.field(pytree_node=False)
    truncate_holding_time: bool = struct.field(pytree_node=False)
    traffic_array: bool = struct.field(pytree_node=False)
    pack_path_bits: bool = struct.field(pytree_node=False)
    relative_arrival_times: bool = struct.field(pytree_node=False)

EnvState

Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

Parameters:

Name	Type	Description	Default
current_time	Scalar	Current time in environment	required
holding_time	Scalar	Holding time of current request	required
total_timesteps	Scalar	Total timesteps in environment	required
total_requests	Scalar	Total requests in environment	required
graph	GraphsTuple	Graph tuple representing network state	required
full_link_slot_mask	Array	Action mask for link slot action (including if slot actions are aggregated)	required
accepted_services	Array	Number of accepted services	required
accepted_bitrate	Array	Accepted bitrate	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvState:
    """Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

    Args:
        current_time (chex.Scalar): Current time in environment
        holding_time (chex.Scalar): Holding time of current request
        total_timesteps (chex.Scalar): Total timesteps in environment
        total_requests (chex.Scalar): Total requests in environment
        graph (jraph.GraphsTuple): Graph tuple representing network state
        full_link_slot_mask (chex.Array): Action mask for link slot action (including if slot actions are aggregated)
        accepted_services (chex.Array): Number of accepted services
        accepted_bitrate (chex.Array): Accepted bitrate
        """
    current_time: chex.Scalar
    holding_time: chex.Scalar
    arrival_time: chex.Scalar
    total_timesteps: chex.Scalar
    total_requests: chex.Scalar
    graph: jraph.GraphsTuple
    full_link_slot_mask: chex.Array
    accepted_services: chex.Array
    accepted_bitrate: chex.Array
    total_bitrate: chex.Array
    list_of_requests: chex.Array

RMSAGNModelEnv

Bases: RSAEnv

RMSA + GNN model environment.

Source code in xlron/environments/gn_model/rmsa_gn_model.py

class RMSAGNModelEnv(RSAEnv):
    """RMSA + GNN model environment."""
    def __init__(
            self,
            key: chex.PRNGKey,
            params: RSAGNModelEnvParams,
            traffic_matrix: chex.Array = None,
            launch_power_array: chex.Array = None,
            list_of_requests: chex.Array = None,
            laplacian_matrix: chex.Array = None,
    ):
        """Initialise the environment state and set as initial state.

        Args:
            key: PRNG key
            params: Environment parameters
            traffic_matrix (optional): Traffic matrix
            launch_power_array (optional): Launch power array

        Returns:
            None
        """
        super().__init__(key, params, traffic_matrix=traffic_matrix, list_of_requests=list_of_requests, laplacian_matrix=laplacian_matrix)
        state = RMSAGNModelEnvState(
            current_time=0,
            arrival_time=0,
            holding_time=0,
            total_timesteps=0,
            total_requests=-1,
            link_slot_array=init_link_slot_array(params),
            link_slot_departure_array=init_link_slot_departure_array(params),
            request_array=init_rsa_request_array(),
            link_slot_mask=init_link_slot_mask(params, agg=params.aggregate_slots),
            traffic_matrix=traffic_matrix if traffic_matrix is not None else init_traffic_matrix(key, params),
            graph=None,
            full_link_slot_mask=init_link_slot_mask(params),
            accepted_services=0,
            accepted_bitrate=0.,
            total_bitrate=0.,
            list_of_requests=list_of_requests,
            link_snr_array=init_link_snr_array(params),
            path_index_array=init_path_index_array(params),
            path_index_array_prev=init_path_index_array(params),
            channel_centre_bw_array=init_channel_centre_bw_array(params),
            channel_power_array=init_channel_power_array(params),
            modulation_format_index_array=init_modulation_format_index_array(params),
            channel_centre_bw_array_prev=init_channel_centre_bw_array(params),
            channel_power_array_prev=init_channel_power_array(params),
            modulation_format_index_array_prev=init_modulation_format_index_array(params),
            launch_power_array=launch_power_array,
            mod_format_mask=init_mod_format_mask(params),
        )
        self.initial_state = state.replace(graph=init_graph_tuple(state, params, laplacian_matrix))

    @partial(jax.jit, static_argnums=(0, 2,))
    def action_mask(self, state: RSAEnvState, params: RSAEnvParams) -> RSAEnvState:
        """Returns mask of valid actions.

        Args:
            state: Environment state
            params: Environment parameters

        Returns:
            state: Environment state with action mask
        """
        state = mask_slots_rmsa_gn_model(state, params, state.request_array)
        return state

    @partial(jax.jit, static_argnums=(0, 2,))
    def get_obs(self, state: RMSAGNModelEnvState, params: RMSAGNModelEnvParams) -> chex.Array:
        return get_paths_obs_gn_model(state, params)

    @staticmethod
    def num_actions(params: RSAEnvParams) -> int:
        """Number of actions possible in environment."""
        return 1

    def observation_space(self, params: RSAEnvParams):
        """Observation space of the environment."""
        return spaces.Discrete(
            3 +  # Request array
            1 +  # Holding time
            7 * params.k_paths
            # Path stats:
            # Mean free block size
            # Free slots
            # Path length (100 km)
            # Path length (hops)
            # Number of connections on path
            # Mean power of connection on path
            # Mean SNR of connection on path
        )

__init__(key, params, traffic_matrix=None, launch_power_array=None, list_of_requests=None, laplacian_matrix=None)

Initialise the environment state and set as initial state.

Parameters:

Name	Type	Description	Default
key	PRNGKey	PRNG key	required
params	RSAGNModelEnvParams	Environment parameters	required
traffic_matrix	optional	Traffic matrix	None
launch_power_array	optional	Launch power array	None

Returns:

Type	Description
	None

Source code in xlron/environments/gn_model/rmsa_gn_model.py

def __init__(
        self,
        key: chex.PRNGKey,
        params: RSAGNModelEnvParams,
        traffic_matrix: chex.Array = None,
        launch_power_array: chex.Array = None,
        list_of_requests: chex.Array = None,
        laplacian_matrix: chex.Array = None,
):
    """Initialise the environment state and set as initial state.

    Args:
        key: PRNG key
        params: Environment parameters
        traffic_matrix (optional): Traffic matrix
        launch_power_array (optional): Launch power array

    Returns:
        None
    """
    super().__init__(key, params, traffic_matrix=traffic_matrix, list_of_requests=list_of_requests, laplacian_matrix=laplacian_matrix)
    state = RMSAGNModelEnvState(
        current_time=0,
        arrival_time=0,
        holding_time=0,
        total_timesteps=0,
        total_requests=-1,
        link_slot_array=init_link_slot_array(params),
        link_slot_departure_array=init_link_slot_departure_array(params),
        request_array=init_rsa_request_array(),
        link_slot_mask=init_link_slot_mask(params, agg=params.aggregate_slots),
        traffic_matrix=traffic_matrix if traffic_matrix is not None else init_traffic_matrix(key, params),
        graph=None,
        full_link_slot_mask=init_link_slot_mask(params),
        accepted_services=0,
        accepted_bitrate=0.,
        total_bitrate=0.,
        list_of_requests=list_of_requests,
        link_snr_array=init_link_snr_array(params),
        path_index_array=init_path_index_array(params),
        path_index_array_prev=init_path_index_array(params),
        channel_centre_bw_array=init_channel_centre_bw_array(params),
        channel_power_array=init_channel_power_array(params),
        modulation_format_index_array=init_modulation_format_index_array(params),
        channel_centre_bw_array_prev=init_channel_centre_bw_array(params),
        channel_power_array_prev=init_channel_power_array(params),
        modulation_format_index_array_prev=init_modulation_format_index_array(params),
        launch_power_array=launch_power_array,
        mod_format_mask=init_mod_format_mask(params),
    )
    self.initial_state = state.replace(graph=init_graph_tuple(state, params, laplacian_matrix))

action_mask(state, params)

Returns mask of valid actions.

Parameters:

Name	Type	Description	Default
state	RSAEnvState	Environment state	required
params	RSAEnvParams	Environment parameters	required

Returns:

Name	Type	Description
state	RSAEnvState	Environment state with action mask

Source code in xlron/environments/gn_model/rmsa_gn_model.py

@partial(jax.jit, static_argnums=(0, 2,))
def action_mask(self, state: RSAEnvState, params: RSAEnvParams) -> RSAEnvState:
    """Returns mask of valid actions.

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        state: Environment state with action mask
    """
    state = mask_slots_rmsa_gn_model(state, params, state.request_array)
    return state

num_actions(params) staticmethod

Number of actions possible in environment.

Source code in xlron/environments/gn_model/rmsa_gn_model.py

@staticmethod
def num_actions(params: RSAEnvParams) -> int:
    """Number of actions possible in environment."""
    return 1

observation_space(params)

Observation space of the environment.

Source code in xlron/environments/gn_model/rmsa_gn_model.py

def observation_space(self, params: RSAEnvParams):
    """Observation space of the environment."""
    return spaces.Discrete(
        3 +  # Request array
        1 +  # Holding time
        7 * params.k_paths
        # Path stats:
        # Mean free block size
        # Free slots
        # Path length (100 km)
        # Path length (hops)
        # Number of connections on path
        # Mean power of connection on path
        # Mean SNR of connection on path
    )

RMSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulations_array: chex.Array = struct.field(pytree_node=False)

RMSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulation_format_index_array: chex.Array  # Modulation format index for each active connection
    modulation_format_index_array_prev: chex.Array  # Modulation format index for each active connection in previous timestep
    mod_format_mask: chex.Array  # Modulation format mask

RSAGNModelEnv

Bases: RSAEnv

RMSA + GNN model environment.

Source code in xlron/environments/gn_model/rsa_gn_model.py

class RSAGNModelEnv(RSAEnv):
    """RMSA + GNN model environment."""
    def __init__(
            self,
            key: chex.PRNGKey,
            params: RSAGNModelEnvParams,
            traffic_matrix: chex.Array = None,
            launch_power_array: chex.Array = None,
            list_of_requests: chex.Array = None,
            laplacian_matrix: chex.Array = None,
    ):
        """Initialise the environment state and set as initial state.

        Args:
            key: PRNG key
            params: Environment parameters
            traffic_matrix (optional): Traffic matrix
            launch_power_array (optional): Launch power array

        Returns:
            None
        """
        super().__init__(key, params, traffic_matrix=traffic_matrix, list_of_requests=list_of_requests, laplacian_matrix=laplacian_matrix)
        state = RSAGNModelEnvState(
            current_time=0,
            arrival_time=0,
            holding_time=0,
            total_timesteps=0,
            total_requests=-1,
            link_slot_array=init_link_slot_array(params),
            link_slot_departure_array=init_link_slot_departure_array(params),
            request_array=init_rsa_request_array(),
            link_slot_mask=init_link_slot_mask(params, agg=params.aggregate_slots),
            traffic_matrix=traffic_matrix if traffic_matrix is not None else init_traffic_matrix(key, params),
            graph=None,
            full_link_slot_mask=init_link_slot_mask(params),
            accepted_services=0,
            accepted_bitrate=0.,
            total_bitrate=0.,
            list_of_requests=list_of_requests,
            link_snr_array=init_link_snr_array(params),
            path_index_array=init_path_index_array(params),
            path_index_array_prev=init_path_index_array(params),
            channel_centre_bw_array=init_channel_centre_bw_array(params),
            channel_power_array=init_channel_power_array(params),
            channel_centre_bw_array_prev=init_channel_centre_bw_array(params),
            channel_power_array_prev=init_channel_power_array(params),
            launch_power_array=launch_power_array,
            active_lightpaths_array=init_active_lightpaths_array(params),
            active_lightpaths_array_departure=init_active_lightpaths_array_departure(params),
            throughput=jnp.array(0., dtype=init_link_snr_array(params).dtype),
        )
        self.initial_state = state.replace(graph=init_graph_tuple(state, params, laplacian_matrix))


    @partial(jax.jit, static_argnums=(0, 2,))
    def reset(
        self, key: chex.PRNGKey, params: Optional[RSAEnvParams] = None, state: Optional[RSAEnvState] = None
    ) -> Tuple[chex.Array, RSAEnvState]:
        """Reset the environment and log the total throughput at episode end.
        """
        throughput_condition = params.monitor_active_lightpaths and state is not None
        throughput = calculate_throughput_from_active_lightpaths(state, params) if (
            throughput_condition) else jnp.array(0.)
        obs, state = super().reset(key, params, state)
        state = state.replace(throughput=throughput)
        jax.debug.print("resetting env, throughput: {}", throughput, ordered=True)
        return obs, state

    @partial(jax.jit, static_argnums=(0, 2,))
    def get_obs(self, state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> chex.Array:
        # Monitoring active lightpaths is used to track total throughput at end of episode
        if params.monitor_active_lightpaths:
           return jnp.array(0)
        return get_paths_obs_gn_model(state, params)

    @staticmethod
    def num_actions(params: RSAEnvParams) -> int:
        """Number of actions possible in environment."""
        return 1

    def observation_space(self, params: RSAEnvParams):
        """Observation space of the environment."""
        return spaces.Discrete(
            3 +  # Request array
            1 +  # Holding time
            7 * params.k_paths
            # Path stats:
            # Mean free block size
            # Free slots
            # Path length (100 km)
            # Path length (hops)
            # Number of connections on path
            # Mean power of connection on path
            # Mean SNR of connection on path
        )

__init__(key, params, traffic_matrix=None, launch_power_array=None, list_of_requests=None, laplacian_matrix=None)

Initialise the environment state and set as initial state.

Parameters:

Name	Type	Description	Default
key	PRNGKey	PRNG key	required
params	RSAGNModelEnvParams	Environment parameters	required
traffic_matrix	optional	Traffic matrix	None
launch_power_array	optional	Launch power array	None

Returns:

Type	Description
	None

Source code in xlron/environments/gn_model/rsa_gn_model.py

def __init__(
        self,
        key: chex.PRNGKey,
        params: RSAGNModelEnvParams,
        traffic_matrix: chex.Array = None,
        launch_power_array: chex.Array = None,
        list_of_requests: chex.Array = None,
        laplacian_matrix: chex.Array = None,
):
    """Initialise the environment state and set as initial state.

    Args:
        key: PRNG key
        params: Environment parameters
        traffic_matrix (optional): Traffic matrix
        launch_power_array (optional): Launch power array

    Returns:
        None
    """
    super().__init__(key, params, traffic_matrix=traffic_matrix, list_of_requests=list_of_requests, laplacian_matrix=laplacian_matrix)
    state = RSAGNModelEnvState(
        current_time=0,
        arrival_time=0,
        holding_time=0,
        total_timesteps=0,
        total_requests=-1,
        link_slot_array=init_link_slot_array(params),
        link_slot_departure_array=init_link_slot_departure_array(params),
        request_array=init_rsa_request_array(),
        link_slot_mask=init_link_slot_mask(params, agg=params.aggregate_slots),
        traffic_matrix=traffic_matrix if traffic_matrix is not None else init_traffic_matrix(key, params),
        graph=None,
        full_link_slot_mask=init_link_slot_mask(params),
        accepted_services=0,
        accepted_bitrate=0.,
        total_bitrate=0.,
        list_of_requests=list_of_requests,
        link_snr_array=init_link_snr_array(params),
        path_index_array=init_path_index_array(params),
        path_index_array_prev=init_path_index_array(params),
        channel_centre_bw_array=init_channel_centre_bw_array(params),
        channel_power_array=init_channel_power_array(params),
        channel_centre_bw_array_prev=init_channel_centre_bw_array(params),
        channel_power_array_prev=init_channel_power_array(params),
        launch_power_array=launch_power_array,
        active_lightpaths_array=init_active_lightpaths_array(params),
        active_lightpaths_array_departure=init_active_lightpaths_array_departure(params),
        throughput=jnp.array(0., dtype=init_link_snr_array(params).dtype),
    )
    self.initial_state = state.replace(graph=init_graph_tuple(state, params, laplacian_matrix))

num_actions(params) staticmethod

Number of actions possible in environment.

Source code in xlron/environments/gn_model/rsa_gn_model.py

@staticmethod
def num_actions(params: RSAEnvParams) -> int:
    """Number of actions possible in environment."""
    return 1

observation_space(params)

Observation space of the environment.

Source code in xlron/environments/gn_model/rsa_gn_model.py

def observation_space(self, params: RSAEnvParams):
    """Observation space of the environment."""
    return spaces.Discrete(
        3 +  # Request array
        1 +  # Holding time
        7 * params.k_paths
        # Path stats:
        # Mean free block size
        # Free slots
        # Path length (100 km)
        # Path length (hops)
        # Number of connections on path
        # Mean power of connection on path
        # Mean SNR of connection on path
    )

reset(key, params=None, state=None)

Reset the environment and log the total throughput at episode end.

Source code in xlron/environments/gn_model/rsa_gn_model.py

@partial(jax.jit, static_argnums=(0, 2,))
def reset(
    self, key: chex.PRNGKey, params: Optional[RSAEnvParams] = None, state: Optional[RSAEnvState] = None
) -> Tuple[chex.Array, RSAEnvState]:
    """Reset the environment and log the total throughput at episode end.
    """
    throughput_condition = params.monitor_active_lightpaths and state is not None
    throughput = calculate_throughput_from_active_lightpaths(state, params) if (
        throughput_condition) else jnp.array(0.)
    obs, state = super().reset(key, params, state)
    state = state.replace(throughput=throughput)
    jax.debug.print("resetting env, throughput: {}", throughput, ordered=True)
    return obs, state

RSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RSA with GN model.
    """
    min_snr: chex.Scalar = struct.field(pytree_node=False)
    fec_threshold: chex.Scalar = struct.field(pytree_node=False)

RSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    active_lightpaths_array: chex.Array  # Active lightpath array. 1 x M array. Each value is a lightpath index. Used to calculate total throughput.
    active_lightpaths_array_departure: chex.Array  # Active lightpath array departure time.
    throughput: chex.Array  # Current network throughput

---

Dataclasses

DeepRMSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for DeepRMSA.

Parameters:

Name	Type	Description	Default
path_stats	Array	Path stats array containing	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class DeepRMSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for DeepRMSA.

    Args:
        path_stats (chex.Array): Path stats array containing
        1. Required slots on path
        2. Total available slots on path
        3. Size of 1st free spectrum block
        4. Avg. free block size
    """
    path_stats: chex.Array

EnvParams

Dataclass to hold environment parameters. Parameters are immutable.

Parameters:

Name	Type	Description	Default
max_requests	Scalar	Maximum number of requests in an episode	required
incremental_loading	Scalar	Incremental increase in traffic load (non-expiring requests)	required
end_first_blocking	Scalar	End episode on first blocking event	required
continuous_operation	Scalar	If True, do not reset the environment at the end of an episode	required
edges	Array	Two column array defining source-dest node-pair edges of the graph	required
slot_size	Scalar	Spectral width of frequency slot in GHz	required
consider_modulation_format	Scalar	If True, consider modulation format to determine required slots	required
link_length_array	Array	Array of link lengths	required
aggregate_slots	Scalar	Number of slots to aggregate into a single action (First-Fit with aggregation)	required
guardband	Scalar	Guard band in slots	required
directed_graph	bool	Whether graph is directed (one fibre per link per transmission direction)	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvParams:
    """Dataclass to hold environment parameters. Parameters are immutable.

    Args:
        max_requests (chex.Scalar): Maximum number of requests in an episode
        incremental_loading (chex.Scalar): Incremental increase in traffic load (non-expiring requests)
        end_first_blocking (chex.Scalar): End episode on first blocking event
        continuous_operation (chex.Scalar): If True, do not reset the environment at the end of an episode
        edges (chex.Array): Two column array defining source-dest node-pair edges of the graph
        slot_size (chex.Scalar): Spectral width of frequency slot in GHz
        consider_modulation_format (chex.Scalar): If True, consider modulation format to determine required slots
        link_length_array (chex.Array): Array of link lengths
        aggregate_slots (chex.Scalar): Number of slots to aggregate into a single action (First-Fit with aggregation)
        guardband (chex.Scalar): Guard band in slots
        directed_graph (bool): Whether graph is directed (one fibre per link per transmission direction)
    """
    max_requests: chex.Scalar = struct.field(pytree_node=False)
    incremental_loading: chex.Scalar = struct.field(pytree_node=False)
    end_first_blocking: chex.Scalar = struct.field(pytree_node=False)
    continuous_operation: chex.Scalar = struct.field(pytree_node=False)
    edges: chex.Array = struct.field(pytree_node=False)
    slot_size: chex.Scalar = struct.field(pytree_node=False)
    consider_modulation_format: chex.Scalar = struct.field(pytree_node=False)
    link_length_array: chex.Array = struct.field(pytree_node=False)
    aggregate_slots: chex.Scalar = struct.field(pytree_node=False)
    guardband: chex.Scalar = struct.field(pytree_node=False)
    directed_graph: bool = struct.field(pytree_node=False)
    maximise_throughput: bool = struct.field(pytree_node=False)
    reward_type: str = struct.field(pytree_node=False)
    values_bw: chex.Array = struct.field(pytree_node=False)
    truncate_holding_time: bool = struct.field(pytree_node=False)
    traffic_array: bool = struct.field(pytree_node=False)
    pack_path_bits: bool = struct.field(pytree_node=False)
    relative_arrival_times: bool = struct.field(pytree_node=False)

EnvState

Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

Parameters:

Name	Type	Description	Default
current_time	Scalar	Current time in environment	required
holding_time	Scalar	Holding time of current request	required
total_timesteps	Scalar	Total timesteps in environment	required
total_requests	Scalar	Total requests in environment	required
graph	GraphsTuple	Graph tuple representing network state	required
full_link_slot_mask	Array	Action mask for link slot action (including if slot actions are aggregated)	required
accepted_services	Array	Number of accepted services	required
accepted_bitrate	Array	Accepted bitrate	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvState:
    """Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

    Args:
        current_time (chex.Scalar): Current time in environment
        holding_time (chex.Scalar): Holding time of current request
        total_timesteps (chex.Scalar): Total timesteps in environment
        total_requests (chex.Scalar): Total requests in environment
        graph (jraph.GraphsTuple): Graph tuple representing network state
        full_link_slot_mask (chex.Array): Action mask for link slot action (including if slot actions are aggregated)
        accepted_services (chex.Array): Number of accepted services
        accepted_bitrate (chex.Array): Accepted bitrate
        """
    current_time: chex.Scalar
    holding_time: chex.Scalar
    arrival_time: chex.Scalar
    total_timesteps: chex.Scalar
    total_requests: chex.Scalar
    graph: jraph.GraphsTuple
    full_link_slot_mask: chex.Array
    accepted_services: chex.Array
    accepted_bitrate: chex.Array
    total_bitrate: chex.Array
    list_of_requests: chex.Array

GNModelEnvParams

Bases: RSAEnvParams

Dataclass to hold environment state for GN model environments.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvParams(RSAEnvParams):
    """Dataclass to hold environment state for GN model environments.
    """
    ref_lambda: chex.Scalar = struct.field(pytree_node=False)
    max_spans: chex.Scalar = struct.field(pytree_node=False)
    max_span_length: chex.Scalar = struct.field(pytree_node=False)
    nonlinear_coeff: chex.Scalar = struct.field(pytree_node=False)
    raman_gain_slope: chex.Scalar = struct.field(pytree_node=False)
    attenuation: chex.Scalar = struct.field(pytree_node=False)
    attenuation_bar: chex.Scalar = struct.field(pytree_node=False)
    dispersion_coeff: chex.Scalar = struct.field(pytree_node=False)
    dispersion_slope: chex.Scalar = struct.field(pytree_node=False)
    transceiver_snr: chex.Array = struct.field(pytree_node=False)
    amplifier_noise_figure: chex.Array = struct.field(pytree_node=False)
    coherent: bool = struct.field(pytree_node=False)
    num_roadms: chex.Scalar = struct.field(pytree_node=False)
    roadm_loss: chex.Scalar = struct.field(pytree_node=False)
    num_spans: chex.Scalar = struct.field(pytree_node=False)
    launch_power_type: chex.Scalar = struct.field(pytree_node=False)
    snr_margin: chex.Scalar = struct.field(pytree_node=False)
    max_snr: chex.Scalar = struct.field(pytree_node=False)
    max_power: chex.Scalar = struct.field(pytree_node=False)
    min_power: chex.Scalar = struct.field(pytree_node=False)
    step_power: chex.Scalar = struct.field(pytree_node=False)
    last_fit: bool = struct.field(pytree_node=False)
    default_launch_power: chex.Scalar = struct.field(pytree_node=False)
    mod_format_correction: bool = struct.field(pytree_node=False)
    monitor_active_lightpaths: bool = struct.field(pytree_node=False)  # Monitor active lightpaths for throughput calculation
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)
    uniform_spans: bool = struct.field(pytree_node=False)

GNModelEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvState(RSAEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    link_snr_array: chex.Array  # Available SNR on each link
    channel_centre_bw_array: chex.Array  # Channel centre bandwidth for each active connection
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots (used for lightpath SNR calculation)
    channel_power_array: chex.Array  # Channel power for each active connection
    channel_centre_bw_array_prev: chex.Array  # Channel centre bandwidth for each active connection in previous timestep
    path_index_array_prev: chex.Array  # Contains indices of lightpaths in use on slots in previous timestep
    channel_power_array_prev: chex.Array  # Channel power for each active connection in previous timestep
    launch_power_array: chex.Array  # Launch power array

LogEnvState

Dataclass to hold environment state for logging.

Parameters:

Name	Type	Description	Default
env_state	EnvState	Environment state	required
lengths	Scalar	Lengths	required
returns	Scalar	Returns	required
cum_returns	Scalar	Cumulative returns	required
episode_lengths	Scalar	Episode lengths	required
episode_returns	Scalar	Episode returns	required
accepted_services	Scalar	Accepted services	required
accepted_bitrate	Scalar	Accepted bitrate	required
done	Scalar	Done	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class LogEnvState:
    """Dataclass to hold environment state for logging.

    Args:
        env_state (EnvState): Environment state
        lengths (chex.Scalar): Lengths
        returns (chex.Scalar): Returns
        cum_returns (chex.Scalar): Cumulative returns
        episode_lengths (chex.Scalar): Episode lengths
        episode_returns (chex.Scalar): Episode returns
        accepted_services (chex.Scalar): Accepted services
        accepted_bitrate (chex.Scalar): Accepted bitrate
        done (chex.Scalar): Done
    """
    env_state: EnvState
    lengths: float
    returns: float
    cum_returns: float
    accepted_services: int
    accepted_bitrate: float
    total_bitrate: float
    utilisation: float
    done: bool

MultiBandRSAEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_start: chex.Scalar = struct.field(pytree_node=False)
    gap_width: chex.Scalar = struct.field(pytree_node=False)

MultiBandRSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RMSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulations_array: chex.Array = struct.field(pytree_node=False)

RMSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulation_format_index_array: chex.Array  # Modulation format index for each active connection
    modulation_format_index_array_prev: chex.Array  # Modulation format index for each active connection in previous timestep
    mod_format_mask: chex.Array  # Modulation format mask

RSAEnvParams

Bases: EnvParams

Dataclass to hold environment parameters for RSA.

Parameters:

Name	Type	Description	Default
num_nodes	Scalar	Number of nodes	required
num_links	Scalar	Number of links	required
link_resources	Scalar	Number of link resources	required
k_paths	Scalar	Number of paths	required
mean_service_holding_time	Scalar	Mean service holding time	required
load	Scalar	Load	required
arrival_rate	Scalar	Arrival rate	required
path_link_array	Array	Path link array	required
random_traffic	bool	Random traffic matrix for RSA on each reset (else uniform or custom)	required
max_slots	Scalar	Maximum number of slots	required
path_se_array	Array	Path spectral efficiency array	required
deterministic_requests	bool	If True, use deterministic requests	required
multiple_topologies	bool	If True, use multiple topologies	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvParams(EnvParams):
    """Dataclass to hold environment parameters for RSA.

    Args:
        num_nodes (chex.Scalar): Number of nodes
        num_links (chex.Scalar): Number of links
        link_resources (chex.Scalar): Number of link resources
        k_paths (chex.Scalar): Number of paths
        mean_service_holding_time (chex.Scalar): Mean service holding time
        load (chex.Scalar): Load
        arrival_rate (chex.Scalar): Arrival rate
        path_link_array (chex.Array): Path link array
        random_traffic (bool): Random traffic matrix for RSA on each reset (else uniform or custom)
        max_slots (chex.Scalar): Maximum number of slots
        path_se_array (chex.Array): Path spectral efficiency array
        deterministic_requests (bool): If True, use deterministic requests
        multiple_topologies (bool): If True, use multiple topologies
    """
    num_nodes: chex.Scalar = struct.field(pytree_node=False)
    num_links: chex.Scalar = struct.field(pytree_node=False)
    link_resources: chex.Scalar = struct.field(pytree_node=False)
    k_paths: chex.Scalar = struct.field(pytree_node=False)
    mean_service_holding_time: chex.Scalar = struct.field(pytree_node=False)
    load: chex.Scalar = struct.field(pytree_node=False)
    arrival_rate: chex.Scalar = struct.field(pytree_node=False)
    path_link_array: chex.Array = struct.field(pytree_node=False)
    random_traffic: bool = struct.field(pytree_node=False)
    max_slots: chex.Scalar = struct.field(pytree_node=False)
    path_se_array: chex.Array = struct.field(pytree_node=False)
    deterministic_requests: bool = struct.field(pytree_node=False)
    multiple_topologies: bool = struct.field(pytree_node=False)
    log_actions: bool = struct.field(pytree_node=False)
    disable_node_features: bool = struct.field(pytree_node=False)

RSAEnvState

Bases: EnvState

Dataclass to hold environment state for RSA.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
request_array	Array	Request array	required
link_slot_departure_array	Array	Link slot departure array	required
link_slot_mask	Array	Link slot mask	required
traffic_matrix	Array	Traffic matrix	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvState(EnvState):
    """Dataclass to hold environment state for RSA.

    Args:
        link_slot_array (chex.Array): Link slot array
        request_array (chex.Array): Request array
        link_slot_departure_array (chex.Array): Link slot departure array
        link_slot_mask (chex.Array): Link slot mask
        traffic_matrix (chex.Array): Traffic matrix
    """
    link_slot_array: chex.Array
    request_array: chex.Array
    link_slot_departure_array: chex.Array
    link_slot_mask: chex.Array
    traffic_matrix: chex.Array

RSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RSA with GN model.
    """
    min_snr: chex.Scalar = struct.field(pytree_node=False)
    fec_threshold: chex.Scalar = struct.field(pytree_node=False)

RSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    active_lightpaths_array: chex.Array  # Active lightpath array. 1 x M array. Each value is a lightpath index. Used to calculate total throughput.
    active_lightpaths_array_departure: chex.Array  # Active lightpath array departure time.
    throughput: chex.Array  # Current network throughput

RSAMultibandEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)

RSAMultibandEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RWALightpathReuseEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RWA with lightpath reuse.

Parameters:

Name	Type	Description	Default
path_index_array	Array	Contains indices of lightpaths in use on slots	required
path_capacity_array	Array	Contains remaining capacity of each lightpath	required
link_capacity_array	Array	Contains remaining capacity of lightpath on each link-slot	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RWALightpathReuseEnvState(RSAEnvState):
    """Dataclass to hold environment state for RWA with lightpath reuse.

    Args:
        path_index_array (chex.Array): Contains indices of lightpaths in use on slots
        path_capacity_array (chex.Array): Contains remaining capacity of each lightpath
        link_capacity_array (chex.Array): Contains remaining capacity of lightpath on each link-slot
    """
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots
    path_capacity_array: chex.Array  # Contains remaining capacity of each lightpath
    link_capacity_array: chex.Array  # Contains remaining capacity of lightpath on each link-slot
    time_since_last_departure: chex.Array  # Time since last departure

VONEEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for VONE.

Parameters:

Name	Type	Description	Default
node_resources	Scalar	Number of node resources	required
max_edges	Scalar	Maximum number of edges	required
min_node_resources	Scalar	Minimum number of node resources	required
max_node_resources	Scalar	Maximum number of node resources	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for VONE.

    Args:
        node_resources (chex.Scalar): Number of node resources
        max_edges (chex.Scalar): Maximum number of edges
        min_node_resources (chex.Scalar): Minimum number of node resources
        max_node_resources (chex.Scalar): Maximum number of node resources
    """
    node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_edges: chex.Scalar = struct.field(pytree_node=False)
    min_node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_node_resources: chex.Scalar = struct.field(pytree_node=False)

VONEEnvState

Bases: RSAEnvState

Dataclass to hold environment state for VONE.

Parameters:

Name	Type	Description	Default
node_capacity_array	Array	Node capacity array	required
node_resource_array	Array	Node resource array	required
node_departure_array	Array	Node departure array	required
action_counter	Array	Action counter	required
action_history	Array	Action history	required
node_mask_s	Array	Node mask for source node	required
node_mask_d	Array	Node mask for destination node	required
virtual_topology_patterns	Array	Virtual topology patterns	required
values_nodes	Array	Values for nodes	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvState(RSAEnvState):
    """Dataclass to hold environment state for VONE.

    Args:
        node_capacity_array (chex.Array): Node capacity array
        node_resource_array (chex.Array): Node resource array
        node_departure_array (chex.Array): Node departure array
        action_counter (chex.Array): Action counter
        action_history (chex.Array): Action history
        node_mask_s (chex.Array): Node mask for source node
        node_mask_d (chex.Array): Node mask for destination node
        virtual_topology_patterns (chex.Array): Virtual topology patterns
        values_nodes (chex.Array): Values for nodes
    """
    node_capacity_array: chex.Array
    node_resource_array: chex.Array
    node_departure_array: chex.Array
    action_counter: chex.Array
    action_history: chex.Array
    node_mask_s: chex.Array
    node_mask_d: chex.Array
    virtual_topology_patterns: chex.Array
    values_nodes: chex.Array

---

Environment wrappers

DeepRMSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for DeepRMSA.

Parameters:

Name	Type	Description	Default
path_stats	Array	Path stats array containing	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class DeepRMSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for DeepRMSA.

    Args:
        path_stats (chex.Array): Path stats array containing
        1. Required slots on path
        2. Total available slots on path
        3. Size of 1st free spectrum block
        4. Avg. free block size
    """
    path_stats: chex.Array

EnvParams

Dataclass to hold environment parameters. Parameters are immutable.

Parameters:

Name	Type	Description	Default
max_requests	Scalar	Maximum number of requests in an episode	required
incremental_loading	Scalar	Incremental increase in traffic load (non-expiring requests)	required
end_first_blocking	Scalar	End episode on first blocking event	required
continuous_operation	Scalar	If True, do not reset the environment at the end of an episode	required
edges	Array	Two column array defining source-dest node-pair edges of the graph	required
slot_size	Scalar	Spectral width of frequency slot in GHz	required
consider_modulation_format	Scalar	If True, consider modulation format to determine required slots	required
link_length_array	Array	Array of link lengths	required
aggregate_slots	Scalar	Number of slots to aggregate into a single action (First-Fit with aggregation)	required
guardband	Scalar	Guard band in slots	required
directed_graph	bool	Whether graph is directed (one fibre per link per transmission direction)	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvParams:
    """Dataclass to hold environment parameters. Parameters are immutable.

    Args:
        max_requests (chex.Scalar): Maximum number of requests in an episode
        incremental_loading (chex.Scalar): Incremental increase in traffic load (non-expiring requests)
        end_first_blocking (chex.Scalar): End episode on first blocking event
        continuous_operation (chex.Scalar): If True, do not reset the environment at the end of an episode
        edges (chex.Array): Two column array defining source-dest node-pair edges of the graph
        slot_size (chex.Scalar): Spectral width of frequency slot in GHz
        consider_modulation_format (chex.Scalar): If True, consider modulation format to determine required slots
        link_length_array (chex.Array): Array of link lengths
        aggregate_slots (chex.Scalar): Number of slots to aggregate into a single action (First-Fit with aggregation)
        guardband (chex.Scalar): Guard band in slots
        directed_graph (bool): Whether graph is directed (one fibre per link per transmission direction)
    """
    max_requests: chex.Scalar = struct.field(pytree_node=False)
    incremental_loading: chex.Scalar = struct.field(pytree_node=False)
    end_first_blocking: chex.Scalar = struct.field(pytree_node=False)
    continuous_operation: chex.Scalar = struct.field(pytree_node=False)
    edges: chex.Array = struct.field(pytree_node=False)
    slot_size: chex.Scalar = struct.field(pytree_node=False)
    consider_modulation_format: chex.Scalar = struct.field(pytree_node=False)
    link_length_array: chex.Array = struct.field(pytree_node=False)
    aggregate_slots: chex.Scalar = struct.field(pytree_node=False)
    guardband: chex.Scalar = struct.field(pytree_node=False)
    directed_graph: bool = struct.field(pytree_node=False)
    maximise_throughput: bool = struct.field(pytree_node=False)
    reward_type: str = struct.field(pytree_node=False)
    values_bw: chex.Array = struct.field(pytree_node=False)
    truncate_holding_time: bool = struct.field(pytree_node=False)
    traffic_array: bool = struct.field(pytree_node=False)
    pack_path_bits: bool = struct.field(pytree_node=False)
    relative_arrival_times: bool = struct.field(pytree_node=False)

EnvState

Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

Parameters:

Name	Type	Description	Default
current_time	Scalar	Current time in environment	required
holding_time	Scalar	Holding time of current request	required
total_timesteps	Scalar	Total timesteps in environment	required
total_requests	Scalar	Total requests in environment	required
graph	GraphsTuple	Graph tuple representing network state	required
full_link_slot_mask	Array	Action mask for link slot action (including if slot actions are aggregated)	required
accepted_services	Array	Number of accepted services	required
accepted_bitrate	Array	Accepted bitrate	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvState:
    """Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

    Args:
        current_time (chex.Scalar): Current time in environment
        holding_time (chex.Scalar): Holding time of current request
        total_timesteps (chex.Scalar): Total timesteps in environment
        total_requests (chex.Scalar): Total requests in environment
        graph (jraph.GraphsTuple): Graph tuple representing network state
        full_link_slot_mask (chex.Array): Action mask for link slot action (including if slot actions are aggregated)
        accepted_services (chex.Array): Number of accepted services
        accepted_bitrate (chex.Array): Accepted bitrate
        """
    current_time: chex.Scalar
    holding_time: chex.Scalar
    arrival_time: chex.Scalar
    total_timesteps: chex.Scalar
    total_requests: chex.Scalar
    graph: jraph.GraphsTuple
    full_link_slot_mask: chex.Array
    accepted_services: chex.Array
    accepted_bitrate: chex.Array
    total_bitrate: chex.Array
    list_of_requests: chex.Array

GNModelEnvParams

Bases: RSAEnvParams

Dataclass to hold environment state for GN model environments.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvParams(RSAEnvParams):
    """Dataclass to hold environment state for GN model environments.
    """
    ref_lambda: chex.Scalar = struct.field(pytree_node=False)
    max_spans: chex.Scalar = struct.field(pytree_node=False)
    max_span_length: chex.Scalar = struct.field(pytree_node=False)
    nonlinear_coeff: chex.Scalar = struct.field(pytree_node=False)
    raman_gain_slope: chex.Scalar = struct.field(pytree_node=False)
    attenuation: chex.Scalar = struct.field(pytree_node=False)
    attenuation_bar: chex.Scalar = struct.field(pytree_node=False)
    dispersion_coeff: chex.Scalar = struct.field(pytree_node=False)
    dispersion_slope: chex.Scalar = struct.field(pytree_node=False)
    transceiver_snr: chex.Array = struct.field(pytree_node=False)
    amplifier_noise_figure: chex.Array = struct.field(pytree_node=False)
    coherent: bool = struct.field(pytree_node=False)
    num_roadms: chex.Scalar = struct.field(pytree_node=False)
    roadm_loss: chex.Scalar = struct.field(pytree_node=False)
    num_spans: chex.Scalar = struct.field(pytree_node=False)
    launch_power_type: chex.Scalar = struct.field(pytree_node=False)
    snr_margin: chex.Scalar = struct.field(pytree_node=False)
    max_snr: chex.Scalar = struct.field(pytree_node=False)
    max_power: chex.Scalar = struct.field(pytree_node=False)
    min_power: chex.Scalar = struct.field(pytree_node=False)
    step_power: chex.Scalar = struct.field(pytree_node=False)
    last_fit: bool = struct.field(pytree_node=False)
    default_launch_power: chex.Scalar = struct.field(pytree_node=False)
    mod_format_correction: bool = struct.field(pytree_node=False)
    monitor_active_lightpaths: bool = struct.field(pytree_node=False)  # Monitor active lightpaths for throughput calculation
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)
    uniform_spans: bool = struct.field(pytree_node=False)

GNModelEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvState(RSAEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    link_snr_array: chex.Array  # Available SNR on each link
    channel_centre_bw_array: chex.Array  # Channel centre bandwidth for each active connection
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots (used for lightpath SNR calculation)
    channel_power_array: chex.Array  # Channel power for each active connection
    channel_centre_bw_array_prev: chex.Array  # Channel centre bandwidth for each active connection in previous timestep
    path_index_array_prev: chex.Array  # Contains indices of lightpaths in use on slots in previous timestep
    channel_power_array_prev: chex.Array  # Channel power for each active connection in previous timestep
    launch_power_array: chex.Array  # Launch power array

HashableArrayWrapper

Bases: Generic[T]

Wrapper for making arrays hashable. In order to access pre-computed data, such as shortest paths between node-pairs or the constituent links of a path, within a jitted function, we need to make the arrays containing this data hashable. By defining this wrapper, we can define a hash method that returns a hash of the array's bytes, thus making the array hashable. From: https://github.com/google/jax/issues/4572#issuecomment-709677518

Source code in xlron/environments/wrappers.py

class HashableArrayWrapper(Generic[T]):
    """Wrapper for making arrays hashable.
    In order to access pre-computed data, such as shortest paths between node-pairs or the constituent links of a path,
    within a jitted function, we need to make the arrays containing this data hashable. By defining this wrapper, we can
    define a __hash__ method that returns a hash of the array's bytes, thus making the array hashable.
    From: https://github.com/google/jax/issues/4572#issuecomment-709677518
    """
    def __init__(self, val: T):
        self.val = val

    def __getattribute__(self, prop):
        if prop == 'val' or prop == "__hash__" or prop == "__eq__":
            return super(HashableArrayWrapper, self).__getattribute__(prop)
        return getattr(self.val, prop)

    def __getitem__(self, key):
        return self.val[key]

    def __setitem__(self, key, val):
        self.val[key] = val

    def __hash__(self):
        return hash(self.val.tobytes())

    def __eq__(self, other):
        if isinstance(other, HashableArrayWrapper):
            return self.__hash__() == other.__hash__()

        f = getattr(self.val, "__eq__")
        return f(self, other)

LogEnvState

Dataclass to hold environment state for logging.

Parameters:

Name	Type	Description	Default
env_state	EnvState	Environment state	required
lengths	Scalar	Lengths	required
returns	Scalar	Returns	required
cum_returns	Scalar	Cumulative returns	required
episode_lengths	Scalar	Episode lengths	required
episode_returns	Scalar	Episode returns	required
accepted_services	Scalar	Accepted services	required
accepted_bitrate	Scalar	Accepted bitrate	required
done	Scalar	Done	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class LogEnvState:
    """Dataclass to hold environment state for logging.

    Args:
        env_state (EnvState): Environment state
        lengths (chex.Scalar): Lengths
        returns (chex.Scalar): Returns
        cum_returns (chex.Scalar): Cumulative returns
        episode_lengths (chex.Scalar): Episode lengths
        episode_returns (chex.Scalar): Episode returns
        accepted_services (chex.Scalar): Accepted services
        accepted_bitrate (chex.Scalar): Accepted bitrate
        done (chex.Scalar): Done
    """
    env_state: EnvState
    lengths: float
    returns: float
    cum_returns: float
    accepted_services: int
    accepted_bitrate: float
    total_bitrate: float
    utilisation: float
    done: bool

LogWrapper

Bases: GymnaxWrapper

Log the episode returns and lengths. Modified from: https://github.com/RobertTLange/gymnax/blob/master/gymnax/wrappers/purerl.py

Source code in xlron/environments/wrappers.py

class LogWrapper(GymnaxWrapper):
    """Log the episode returns and lengths.
    Modified from: https://github.com/RobertTLange/gymnax/blob/master/gymnax/wrappers/purerl.py
    """

    def __init__(self, env: environment.Environment):
        super().__init__(env)

    @partial(jax.jit, static_argnums=(0,))
    def reset(
        self, key: chex.PRNGKey, params: Optional[environment.EnvParams] = None, state: Optional[environment.EnvState] = None
    ) -> Tuple[chex.Array, environment.EnvState]:
        obs, env_state = self._env.reset(key, params, state)
        state = LogEnvState(
            env_state=env_state,
            lengths=0,
            returns=0,
            cum_returns=0,
            accepted_services=0,
            accepted_bitrate=0,
            total_bitrate=0,
            utilisation=0,
            done=False,
        )
        return obs, state

    @partial(jax.jit, static_argnums=(0,))
    def step(
        self,
        key: chex.PRNGKey,
        state: environment.EnvState,
        action: Union[int, float],
        params: Optional[environment.EnvParams] = None,
    ) -> Tuple[chex.Array, environment.EnvState, float, bool, dict]:
        obs, env_state, reward, done, info = self._env.step(
            key, state.env_state, action, params
        )
        state = LogEnvState(
            env_state=env_state,
            lengths=state.lengths * (1 - done) + 1,
            returns=reward,
            cum_returns=state.cum_returns * (1 - done) + reward,
            accepted_services=env_state.accepted_services,
            accepted_bitrate=env_state.accepted_bitrate,
            total_bitrate=env_state.total_bitrate,
            utilisation=jnp.count_nonzero(env_state.link_slot_array) / env_state.link_slot_array.size,
            done=done,
        )
        info["lengths"] = state.lengths
        info["returns"] = state.returns
        info["cum_returns"] = state.cum_returns
        info["accepted_services"] = state.accepted_services
        info["accepted_bitrate"] = state.accepted_bitrate
        info["total_bitrate"] = state.total_bitrate
        info["utilisation"] = state.utilisation
        info["done"] = done
        # First check if we're dealing with RSAGNModelEnvParams
        is_rsa_params = params.__class__.__name__ == "RSAGNModelEnvParams"

        # For RSA params, unpack the action
        if is_rsa_params:
            action, power_action = action
            info["launch_power"] = power_action

        # Now, if we need to log actions OR we have RSA params, compute the common fields
        if is_rsa_params or params.log_actions:
            # Compute common fields
            nodes_sd, dr_request = read_rsa_request(state.env_state.request_array)
            source, dest = nodes_sd
            i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph).astype(
                jnp.int32)
            path_index, slot_index = process_path_action(state.env_state, params, action)

            # Set common info
            info["path_index"] = i + path_index
            info["slot_index"] = slot_index
            info["source"] = source
            info["dest"] = dest
            info["data_rate"] = dr_request[0]

            # RSA-specific throughput info
            if is_rsa_params:
                info["throughput"] = env_state.throughput #* done  # Only != 0 at end of episode

            # Logging-specific info
            if params.log_actions:
                # RSA-specific logging
                if is_rsa_params:
                    path = params.path_link_array.val[path_index.astype(jnp.int32)]
                    info["path_snr"] = get_snr_for_path(path, env_state.link_snr_array, params)[
                        slot_index.astype(jnp.int32)]
                # Common logging fields
                info["arrival_time"] = env_state.current_time[0]
                info["departure_time"] = env_state.current_time[0] + env_state.holding_time[0]
        return obs, state, reward, done, info

    def _tree_flatten(self):
        children = ()  # arrays / dynamic values
        aux_data = {"env": self._env}  # static values
        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children):
        return cls(*children, **aux_data)

MultiBandRSAEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_start: chex.Scalar = struct.field(pytree_node=False)
    gap_width: chex.Scalar = struct.field(pytree_node=False)

MultiBandRSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RMSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulations_array: chex.Array = struct.field(pytree_node=False)

RMSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulation_format_index_array: chex.Array  # Modulation format index for each active connection
    modulation_format_index_array_prev: chex.Array  # Modulation format index for each active connection in previous timestep
    mod_format_mask: chex.Array  # Modulation format mask

RSAEnvParams

Bases: EnvParams

Dataclass to hold environment parameters for RSA.

Parameters:

Name	Type	Description	Default
num_nodes	Scalar	Number of nodes	required
num_links	Scalar	Number of links	required
link_resources	Scalar	Number of link resources	required
k_paths	Scalar	Number of paths	required
mean_service_holding_time	Scalar	Mean service holding time	required
load	Scalar	Load	required
arrival_rate	Scalar	Arrival rate	required
path_link_array	Array	Path link array	required
random_traffic	bool	Random traffic matrix for RSA on each reset (else uniform or custom)	required
max_slots	Scalar	Maximum number of slots	required
path_se_array	Array	Path spectral efficiency array	required
deterministic_requests	bool	If True, use deterministic requests	required
multiple_topologies	bool	If True, use multiple topologies	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvParams(EnvParams):
    """Dataclass to hold environment parameters for RSA.

    Args:
        num_nodes (chex.Scalar): Number of nodes
        num_links (chex.Scalar): Number of links
        link_resources (chex.Scalar): Number of link resources
        k_paths (chex.Scalar): Number of paths
        mean_service_holding_time (chex.Scalar): Mean service holding time
        load (chex.Scalar): Load
        arrival_rate (chex.Scalar): Arrival rate
        path_link_array (chex.Array): Path link array
        random_traffic (bool): Random traffic matrix for RSA on each reset (else uniform or custom)
        max_slots (chex.Scalar): Maximum number of slots
        path_se_array (chex.Array): Path spectral efficiency array
        deterministic_requests (bool): If True, use deterministic requests
        multiple_topologies (bool): If True, use multiple topologies
    """
    num_nodes: chex.Scalar = struct.field(pytree_node=False)
    num_links: chex.Scalar = struct.field(pytree_node=False)
    link_resources: chex.Scalar = struct.field(pytree_node=False)
    k_paths: chex.Scalar = struct.field(pytree_node=False)
    mean_service_holding_time: chex.Scalar = struct.field(pytree_node=False)
    load: chex.Scalar = struct.field(pytree_node=False)
    arrival_rate: chex.Scalar = struct.field(pytree_node=False)
    path_link_array: chex.Array = struct.field(pytree_node=False)
    random_traffic: bool = struct.field(pytree_node=False)
    max_slots: chex.Scalar = struct.field(pytree_node=False)
    path_se_array: chex.Array = struct.field(pytree_node=False)
    deterministic_requests: bool = struct.field(pytree_node=False)
    multiple_topologies: bool = struct.field(pytree_node=False)
    log_actions: bool = struct.field(pytree_node=False)
    disable_node_features: bool = struct.field(pytree_node=False)

RSAEnvState

Bases: EnvState

Dataclass to hold environment state for RSA.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
request_array	Array	Request array	required
link_slot_departure_array	Array	Link slot departure array	required
link_slot_mask	Array	Link slot mask	required
traffic_matrix	Array	Traffic matrix	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvState(EnvState):
    """Dataclass to hold environment state for RSA.

    Args:
        link_slot_array (chex.Array): Link slot array
        request_array (chex.Array): Request array
        link_slot_departure_array (chex.Array): Link slot departure array
        link_slot_mask (chex.Array): Link slot mask
        traffic_matrix (chex.Array): Traffic matrix
    """
    link_slot_array: chex.Array
    request_array: chex.Array
    link_slot_departure_array: chex.Array
    link_slot_mask: chex.Array
    traffic_matrix: chex.Array

RSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RSA with GN model.
    """
    min_snr: chex.Scalar = struct.field(pytree_node=False)
    fec_threshold: chex.Scalar = struct.field(pytree_node=False)

RSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    active_lightpaths_array: chex.Array  # Active lightpath array. 1 x M array. Each value is a lightpath index. Used to calculate total throughput.
    active_lightpaths_array_departure: chex.Array  # Active lightpath array departure time.
    throughput: chex.Array  # Current network throughput

RSAMultibandEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)

RSAMultibandEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RWALightpathReuseEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RWA with lightpath reuse.

Parameters:

Name	Type	Description	Default
path_index_array	Array	Contains indices of lightpaths in use on slots	required
path_capacity_array	Array	Contains remaining capacity of each lightpath	required
link_capacity_array	Array	Contains remaining capacity of lightpath on each link-slot	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RWALightpathReuseEnvState(RSAEnvState):
    """Dataclass to hold environment state for RWA with lightpath reuse.

    Args:
        path_index_array (chex.Array): Contains indices of lightpaths in use on slots
        path_capacity_array (chex.Array): Contains remaining capacity of each lightpath
        link_capacity_array (chex.Array): Contains remaining capacity of lightpath on each link-slot
    """
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots
    path_capacity_array: chex.Array  # Contains remaining capacity of each lightpath
    link_capacity_array: chex.Array  # Contains remaining capacity of lightpath on each link-slot
    time_since_last_departure: chex.Array  # Time since last departure

RolloutWrapper

Wrapper to define batch evaluation for generation parameters. Used for genetic algorithm. From: https://github.com/RobertTLange/gymnax/

Source code in xlron/environments/wrappers.py

class RolloutWrapper:
    """Wrapper to define batch evaluation for generation parameters. Used for genetic algorithm.
    From: https://github.com/RobertTLange/gymnax/
    """
    def __init__(
        self,
        model_forward=None,
        env: environment.Environment = None,
        num_env_steps: Optional[int] = None,
        env_params: EnvParams = None,
    ):
        """Wrapper to define batch evaluation for generation parameters."""
        self.env = env
        # Define the RL environment & network forward function
        self.env_params = env_params
        self.model_forward = model_forward

        if num_env_steps is None:
            self.num_env_steps = self.env_params.max_requests
        else:
            self.num_env_steps = num_env_steps

    @partial(jax.jit, static_argnums=(0, 2))
    def population_rollout(self, rng_eval, policy_params):
        """Reshape parameter vector and evaluate the generation."""
        # Evaluate population of nets on gymnax task - vmap over rng & params
        pop_rollout = jax.vmap(self.batch_rollout, in_axes=(None, 0))
        return pop_rollout(rng_eval, policy_params)

    @partial(jax.jit, static_argnums=(0, 2))
    def batch_rollout(self, rng_eval, policy_params):
        """Evaluate a generation of networks on RL/Supervised/etc. task."""
        # vmap over different MC fitness evaluations for single network
        batch_rollout = jax.vmap(self.single_rollout, in_axes=(0, None))
        return batch_rollout(rng_eval, policy_params)

    @partial(jax.jit, static_argnums=(0, 2))
    def single_rollout(self, rng_input, policy_params):
        """Rollout a pendulum episode with lax.scan."""
        # Reset the environment
        rng_reset, rng_episode = jax.random.split(rng_input)
        obs, state = self.env.reset(rng_reset, self.env_params)

        def policy_step(state_input, tmp):
            """lax.scan compatible step transition in jax env."""
            obs, state, policy_params, rng, cum_reward, valid_mask = state_input
            rng, rng_step, rng_net = jax.random.split(rng, 3)
            if self.model_forward is not None:
                action = self.model_forward(policy_params, obs, rng_net)
            else:
                action = self.env.action_space(self.env_params).sample(rng_net)
            next_obs, next_state, reward, done, _ = self.env.step(
                rng_step, state, action, self.env_params
            )
            new_cum_reward = cum_reward + reward * valid_mask
            new_valid_mask = valid_mask * (1 - done)
            carry = [
                next_obs,
                next_state,
                policy_params,
                rng,
                new_cum_reward,
                new_valid_mask,
            ]
            y = [obs, action, reward, next_obs, done]
            return carry, y

        # Scan over episode step loop
        carry_out, scan_out = jax.lax.scan(
            policy_step,
            [
                obs,
                state,
                policy_params,
                rng_episode,
                jnp.array([0.0]),
                jnp.array([1.0]),
            ],
            (),
            self.num_env_steps,
        )
        # Return the sum of rewards accumulated by agent in episode rollout
        obs, action, reward, next_obs, done = scan_out
        cum_return = carry_out[-2]
        return obs, action, reward, next_obs, done, cum_return

    @property
    def input_shape(self):
        """Get the shape of the observation."""
        rng = jax.random.PRNGKey(0)
        obs, state = self.env.reset(rng, self.env_params)
        return obs.shape

input_shape property

Get the shape of the observation.

__init__(model_forward=None, env=None, num_env_steps=None, env_params=None)

Wrapper to define batch evaluation for generation parameters.

Source code in xlron/environments/wrappers.py

def __init__(
    self,
    model_forward=None,
    env: environment.Environment = None,
    num_env_steps: Optional[int] = None,
    env_params: EnvParams = None,
):
    """Wrapper to define batch evaluation for generation parameters."""
    self.env = env
    # Define the RL environment & network forward function
    self.env_params = env_params
    self.model_forward = model_forward

    if num_env_steps is None:
        self.num_env_steps = self.env_params.max_requests
    else:
        self.num_env_steps = num_env_steps

batch_rollout(rng_eval, policy_params)

Evaluate a generation of networks on RL/Supervised/etc. task.

Source code in xlron/environments/wrappers.py

@partial(jax.jit, static_argnums=(0, 2))
def batch_rollout(self, rng_eval, policy_params):
    """Evaluate a generation of networks on RL/Supervised/etc. task."""
    # vmap over different MC fitness evaluations for single network
    batch_rollout = jax.vmap(self.single_rollout, in_axes=(0, None))
    return batch_rollout(rng_eval, policy_params)

population_rollout(rng_eval, policy_params)

Reshape parameter vector and evaluate the generation.

Source code in xlron/environments/wrappers.py

@partial(jax.jit, static_argnums=(0, 2))
def population_rollout(self, rng_eval, policy_params):
    """Reshape parameter vector and evaluate the generation."""
    # Evaluate population of nets on gymnax task - vmap over rng & params
    pop_rollout = jax.vmap(self.batch_rollout, in_axes=(None, 0))
    return pop_rollout(rng_eval, policy_params)

single_rollout(rng_input, policy_params)

Rollout a pendulum episode with lax.scan.

Source code in xlron/environments/wrappers.py

@partial(jax.jit, static_argnums=(0, 2))
def single_rollout(self, rng_input, policy_params):
    """Rollout a pendulum episode with lax.scan."""
    # Reset the environment
    rng_reset, rng_episode = jax.random.split(rng_input)
    obs, state = self.env.reset(rng_reset, self.env_params)

    def policy_step(state_input, tmp):
        """lax.scan compatible step transition in jax env."""
        obs, state, policy_params, rng, cum_reward, valid_mask = state_input
        rng, rng_step, rng_net = jax.random.split(rng, 3)
        if self.model_forward is not None:
            action = self.model_forward(policy_params, obs, rng_net)
        else:
            action = self.env.action_space(self.env_params).sample(rng_net)
        next_obs, next_state, reward, done, _ = self.env.step(
            rng_step, state, action, self.env_params
        )
        new_cum_reward = cum_reward + reward * valid_mask
        new_valid_mask = valid_mask * (1 - done)
        carry = [
            next_obs,
            next_state,
            policy_params,
            rng,
            new_cum_reward,
            new_valid_mask,
        ]
        y = [obs, action, reward, next_obs, done]
        return carry, y

    # Scan over episode step loop
    carry_out, scan_out = jax.lax.scan(
        policy_step,
        [
            obs,
            state,
            policy_params,
            rng_episode,
            jnp.array([0.0]),
            jnp.array([1.0]),
        ],
        (),
        self.num_env_steps,
    )
    # Return the sum of rewards accumulated by agent in episode rollout
    obs, action, reward, next_obs, done = scan_out
    cum_return = carry_out[-2]
    return obs, action, reward, next_obs, done, cum_return

TimeIt

Context manager for timing execution of code blocks.

Source code in xlron/environments/wrappers.py

class TimeIt:
    """Context manager for timing execution of code blocks."""

    def __init__(self, tag, frames=None):
        self.tag = tag
        self.frames = frames

    def __enter__(self):
        self.start = timeit.default_timer()
        return self

    def __exit__(self, *args):
        self.elapsed_secs = timeit.default_timer() - self.start
        msg = self.tag + (': Elapsed time=%.2fs' % self.elapsed_secs)
        if self.frames:
            msg += ', FPS=%.2e' % (self.frames / self.elapsed_secs)
        print(msg)

VONEEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for VONE.

Parameters:

Name	Type	Description	Default
node_resources	Scalar	Number of node resources	required
max_edges	Scalar	Maximum number of edges	required
min_node_resources	Scalar	Minimum number of node resources	required
max_node_resources	Scalar	Maximum number of node resources	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for VONE.

    Args:
        node_resources (chex.Scalar): Number of node resources
        max_edges (chex.Scalar): Maximum number of edges
        min_node_resources (chex.Scalar): Minimum number of node resources
        max_node_resources (chex.Scalar): Maximum number of node resources
    """
    node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_edges: chex.Scalar = struct.field(pytree_node=False)
    min_node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_node_resources: chex.Scalar = struct.field(pytree_node=False)

VONEEnvState

Bases: RSAEnvState

Dataclass to hold environment state for VONE.

Parameters:

Name	Type	Description	Default
node_capacity_array	Array	Node capacity array	required
node_resource_array	Array	Node resource array	required
node_departure_array	Array	Node departure array	required
action_counter	Array	Action counter	required
action_history	Array	Action history	required
node_mask_s	Array	Node mask for source node	required
node_mask_d	Array	Node mask for destination node	required
virtual_topology_patterns	Array	Virtual topology patterns	required
values_nodes	Array	Values for nodes	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvState(RSAEnvState):
    """Dataclass to hold environment state for VONE.

    Args:
        node_capacity_array (chex.Array): Node capacity array
        node_resource_array (chex.Array): Node resource array
        node_departure_array (chex.Array): Node departure array
        action_counter (chex.Array): Action counter
        action_history (chex.Array): Action history
        node_mask_s (chex.Array): Node mask for source node
        node_mask_d (chex.Array): Node mask for destination node
        virtual_topology_patterns (chex.Array): Virtual topology patterns
        values_nodes (chex.Array): Values for nodes
    """
    node_capacity_array: chex.Array
    node_resource_array: chex.Array
    node_departure_array: chex.Array
    action_counter: chex.Array
    action_history: chex.Array
    node_mask_s: chex.Array
    node_mask_d: chex.Array
    virtual_topology_patterns: chex.Array
    values_nodes: chex.Array

---

Constituent functions of environments

DeepRMSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for DeepRMSA.

Parameters:

Name	Type	Description	Default
path_stats	Array	Path stats array containing	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class DeepRMSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for DeepRMSA.

    Args:
        path_stats (chex.Array): Path stats array containing
        1. Required slots on path
        2. Total available slots on path
        3. Size of 1st free spectrum block
        4. Avg. free block size
    """
    path_stats: chex.Array

EnvParams

Dataclass to hold environment parameters. Parameters are immutable.

Parameters:

Name	Type	Description	Default
max_requests	Scalar	Maximum number of requests in an episode	required
incremental_loading	Scalar	Incremental increase in traffic load (non-expiring requests)	required
end_first_blocking	Scalar	End episode on first blocking event	required
continuous_operation	Scalar	If True, do not reset the environment at the end of an episode	required
edges	Array	Two column array defining source-dest node-pair edges of the graph	required
slot_size	Scalar	Spectral width of frequency slot in GHz	required
consider_modulation_format	Scalar	If True, consider modulation format to determine required slots	required
link_length_array	Array	Array of link lengths	required
aggregate_slots	Scalar	Number of slots to aggregate into a single action (First-Fit with aggregation)	required
guardband	Scalar	Guard band in slots	required
directed_graph	bool	Whether graph is directed (one fibre per link per transmission direction)	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvParams:
    """Dataclass to hold environment parameters. Parameters are immutable.

    Args:
        max_requests (chex.Scalar): Maximum number of requests in an episode
        incremental_loading (chex.Scalar): Incremental increase in traffic load (non-expiring requests)
        end_first_blocking (chex.Scalar): End episode on first blocking event
        continuous_operation (chex.Scalar): If True, do not reset the environment at the end of an episode
        edges (chex.Array): Two column array defining source-dest node-pair edges of the graph
        slot_size (chex.Scalar): Spectral width of frequency slot in GHz
        consider_modulation_format (chex.Scalar): If True, consider modulation format to determine required slots
        link_length_array (chex.Array): Array of link lengths
        aggregate_slots (chex.Scalar): Number of slots to aggregate into a single action (First-Fit with aggregation)
        guardband (chex.Scalar): Guard band in slots
        directed_graph (bool): Whether graph is directed (one fibre per link per transmission direction)
    """
    max_requests: chex.Scalar = struct.field(pytree_node=False)
    incremental_loading: chex.Scalar = struct.field(pytree_node=False)
    end_first_blocking: chex.Scalar = struct.field(pytree_node=False)
    continuous_operation: chex.Scalar = struct.field(pytree_node=False)
    edges: chex.Array = struct.field(pytree_node=False)
    slot_size: chex.Scalar = struct.field(pytree_node=False)
    consider_modulation_format: chex.Scalar = struct.field(pytree_node=False)
    link_length_array: chex.Array = struct.field(pytree_node=False)
    aggregate_slots: chex.Scalar = struct.field(pytree_node=False)
    guardband: chex.Scalar = struct.field(pytree_node=False)
    directed_graph: bool = struct.field(pytree_node=False)
    maximise_throughput: bool = struct.field(pytree_node=False)
    reward_type: str = struct.field(pytree_node=False)
    values_bw: chex.Array = struct.field(pytree_node=False)
    truncate_holding_time: bool = struct.field(pytree_node=False)
    traffic_array: bool = struct.field(pytree_node=False)
    pack_path_bits: bool = struct.field(pytree_node=False)
    relative_arrival_times: bool = struct.field(pytree_node=False)

EnvState

Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

Parameters:

Name	Type	Description	Default
current_time	Scalar	Current time in environment	required
holding_time	Scalar	Holding time of current request	required
total_timesteps	Scalar	Total timesteps in environment	required
total_requests	Scalar	Total requests in environment	required
graph	GraphsTuple	Graph tuple representing network state	required
full_link_slot_mask	Array	Action mask for link slot action (including if slot actions are aggregated)	required
accepted_services	Array	Number of accepted services	required
accepted_bitrate	Array	Accepted bitrate	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvState:
    """Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

    Args:
        current_time (chex.Scalar): Current time in environment
        holding_time (chex.Scalar): Holding time of current request
        total_timesteps (chex.Scalar): Total timesteps in environment
        total_requests (chex.Scalar): Total requests in environment
        graph (jraph.GraphsTuple): Graph tuple representing network state
        full_link_slot_mask (chex.Array): Action mask for link slot action (including if slot actions are aggregated)
        accepted_services (chex.Array): Number of accepted services
        accepted_bitrate (chex.Array): Accepted bitrate
        """
    current_time: chex.Scalar
    holding_time: chex.Scalar
    arrival_time: chex.Scalar
    total_timesteps: chex.Scalar
    total_requests: chex.Scalar
    graph: jraph.GraphsTuple
    full_link_slot_mask: chex.Array
    accepted_services: chex.Array
    accepted_bitrate: chex.Array
    total_bitrate: chex.Array
    list_of_requests: chex.Array

GNModelEnvParams

Bases: RSAEnvParams

Dataclass to hold environment state for GN model environments.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvParams(RSAEnvParams):
    """Dataclass to hold environment state for GN model environments.
    """
    ref_lambda: chex.Scalar = struct.field(pytree_node=False)
    max_spans: chex.Scalar = struct.field(pytree_node=False)
    max_span_length: chex.Scalar = struct.field(pytree_node=False)
    nonlinear_coeff: chex.Scalar = struct.field(pytree_node=False)
    raman_gain_slope: chex.Scalar = struct.field(pytree_node=False)
    attenuation: chex.Scalar = struct.field(pytree_node=False)
    attenuation_bar: chex.Scalar = struct.field(pytree_node=False)
    dispersion_coeff: chex.Scalar = struct.field(pytree_node=False)
    dispersion_slope: chex.Scalar = struct.field(pytree_node=False)
    transceiver_snr: chex.Array = struct.field(pytree_node=False)
    amplifier_noise_figure: chex.Array = struct.field(pytree_node=False)
    coherent: bool = struct.field(pytree_node=False)
    num_roadms: chex.Scalar = struct.field(pytree_node=False)
    roadm_loss: chex.Scalar = struct.field(pytree_node=False)
    num_spans: chex.Scalar = struct.field(pytree_node=False)
    launch_power_type: chex.Scalar = struct.field(pytree_node=False)
    snr_margin: chex.Scalar = struct.field(pytree_node=False)
    max_snr: chex.Scalar = struct.field(pytree_node=False)
    max_power: chex.Scalar = struct.field(pytree_node=False)
    min_power: chex.Scalar = struct.field(pytree_node=False)
    step_power: chex.Scalar = struct.field(pytree_node=False)
    last_fit: bool = struct.field(pytree_node=False)
    default_launch_power: chex.Scalar = struct.field(pytree_node=False)
    mod_format_correction: bool = struct.field(pytree_node=False)
    monitor_active_lightpaths: bool = struct.field(pytree_node=False)  # Monitor active lightpaths for throughput calculation
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)
    uniform_spans: bool = struct.field(pytree_node=False)

GNModelEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvState(RSAEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    link_snr_array: chex.Array  # Available SNR on each link
    channel_centre_bw_array: chex.Array  # Channel centre bandwidth for each active connection
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots (used for lightpath SNR calculation)
    channel_power_array: chex.Array  # Channel power for each active connection
    channel_centre_bw_array_prev: chex.Array  # Channel centre bandwidth for each active connection in previous timestep
    path_index_array_prev: chex.Array  # Contains indices of lightpaths in use on slots in previous timestep
    channel_power_array_prev: chex.Array  # Channel power for each active connection in previous timestep
    launch_power_array: chex.Array  # Launch power array

LogEnvState

Dataclass to hold environment state for logging.

Parameters:

Name	Type	Description	Default
env_state	EnvState	Environment state	required
lengths	Scalar	Lengths	required
returns	Scalar	Returns	required
cum_returns	Scalar	Cumulative returns	required
episode_lengths	Scalar	Episode lengths	required
episode_returns	Scalar	Episode returns	required
accepted_services	Scalar	Accepted services	required
accepted_bitrate	Scalar	Accepted bitrate	required
done	Scalar	Done	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class LogEnvState:
    """Dataclass to hold environment state for logging.

    Args:
        env_state (EnvState): Environment state
        lengths (chex.Scalar): Lengths
        returns (chex.Scalar): Returns
        cum_returns (chex.Scalar): Cumulative returns
        episode_lengths (chex.Scalar): Episode lengths
        episode_returns (chex.Scalar): Episode returns
        accepted_services (chex.Scalar): Accepted services
        accepted_bitrate (chex.Scalar): Accepted bitrate
        done (chex.Scalar): Done
    """
    env_state: EnvState
    lengths: float
    returns: float
    cum_returns: float
    accepted_services: int
    accepted_bitrate: float
    total_bitrate: float
    utilisation: float
    done: bool

MultiBandRSAEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_start: chex.Scalar = struct.field(pytree_node=False)
    gap_width: chex.Scalar = struct.field(pytree_node=False)

MultiBandRSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RMSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulations_array: chex.Array = struct.field(pytree_node=False)

RMSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulation_format_index_array: chex.Array  # Modulation format index for each active connection
    modulation_format_index_array_prev: chex.Array  # Modulation format index for each active connection in previous timestep
    mod_format_mask: chex.Array  # Modulation format mask

RSAEnvParams

Bases: EnvParams

Dataclass to hold environment parameters for RSA.

Parameters:

Name	Type	Description	Default
num_nodes	Scalar	Number of nodes	required
num_links	Scalar	Number of links	required
link_resources	Scalar	Number of link resources	required
k_paths	Scalar	Number of paths	required
mean_service_holding_time	Scalar	Mean service holding time	required
load	Scalar	Load	required
arrival_rate	Scalar	Arrival rate	required
path_link_array	Array	Path link array	required
random_traffic	bool	Random traffic matrix for RSA on each reset (else uniform or custom)	required
max_slots	Scalar	Maximum number of slots	required
path_se_array	Array	Path spectral efficiency array	required
deterministic_requests	bool	If True, use deterministic requests	required
multiple_topologies	bool	If True, use multiple topologies	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvParams(EnvParams):
    """Dataclass to hold environment parameters for RSA.

    Args:
        num_nodes (chex.Scalar): Number of nodes
        num_links (chex.Scalar): Number of links
        link_resources (chex.Scalar): Number of link resources
        k_paths (chex.Scalar): Number of paths
        mean_service_holding_time (chex.Scalar): Mean service holding time
        load (chex.Scalar): Load
        arrival_rate (chex.Scalar): Arrival rate
        path_link_array (chex.Array): Path link array
        random_traffic (bool): Random traffic matrix for RSA on each reset (else uniform or custom)
        max_slots (chex.Scalar): Maximum number of slots
        path_se_array (chex.Array): Path spectral efficiency array
        deterministic_requests (bool): If True, use deterministic requests
        multiple_topologies (bool): If True, use multiple topologies
    """
    num_nodes: chex.Scalar = struct.field(pytree_node=False)
    num_links: chex.Scalar = struct.field(pytree_node=False)
    link_resources: chex.Scalar = struct.field(pytree_node=False)
    k_paths: chex.Scalar = struct.field(pytree_node=False)
    mean_service_holding_time: chex.Scalar = struct.field(pytree_node=False)
    load: chex.Scalar = struct.field(pytree_node=False)
    arrival_rate: chex.Scalar = struct.field(pytree_node=False)
    path_link_array: chex.Array = struct.field(pytree_node=False)
    random_traffic: bool = struct.field(pytree_node=False)
    max_slots: chex.Scalar = struct.field(pytree_node=False)
    path_se_array: chex.Array = struct.field(pytree_node=False)
    deterministic_requests: bool = struct.field(pytree_node=False)
    multiple_topologies: bool = struct.field(pytree_node=False)
    log_actions: bool = struct.field(pytree_node=False)
    disable_node_features: bool = struct.field(pytree_node=False)

RSAEnvState

Bases: EnvState

Dataclass to hold environment state for RSA.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
request_array	Array	Request array	required
link_slot_departure_array	Array	Link slot departure array	required
link_slot_mask	Array	Link slot mask	required
traffic_matrix	Array	Traffic matrix	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvState(EnvState):
    """Dataclass to hold environment state for RSA.

    Args:
        link_slot_array (chex.Array): Link slot array
        request_array (chex.Array): Request array
        link_slot_departure_array (chex.Array): Link slot departure array
        link_slot_mask (chex.Array): Link slot mask
        traffic_matrix (chex.Array): Traffic matrix
    """
    link_slot_array: chex.Array
    request_array: chex.Array
    link_slot_departure_array: chex.Array
    link_slot_mask: chex.Array
    traffic_matrix: chex.Array

RSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RSA with GN model.
    """
    min_snr: chex.Scalar = struct.field(pytree_node=False)
    fec_threshold: chex.Scalar = struct.field(pytree_node=False)

RSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    active_lightpaths_array: chex.Array  # Active lightpath array. 1 x M array. Each value is a lightpath index. Used to calculate total throughput.
    active_lightpaths_array_departure: chex.Array  # Active lightpath array departure time.
    throughput: chex.Array  # Current network throughput

RSAMultibandEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)

RSAMultibandEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RWALightpathReuseEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RWA with lightpath reuse.

Parameters:

Name	Type	Description	Default
path_index_array	Array	Contains indices of lightpaths in use on slots	required
path_capacity_array	Array	Contains remaining capacity of each lightpath	required
link_capacity_array	Array	Contains remaining capacity of lightpath on each link-slot	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RWALightpathReuseEnvState(RSAEnvState):
    """Dataclass to hold environment state for RWA with lightpath reuse.

    Args:
        path_index_array (chex.Array): Contains indices of lightpaths in use on slots
        path_capacity_array (chex.Array): Contains remaining capacity of each lightpath
        link_capacity_array (chex.Array): Contains remaining capacity of lightpath on each link-slot
    """
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots
    path_capacity_array: chex.Array  # Contains remaining capacity of each lightpath
    link_capacity_array: chex.Array  # Contains remaining capacity of lightpath on each link-slot
    time_since_last_departure: chex.Array  # Time since last departure

VONEEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for VONE.

Parameters:

Name	Type	Description	Default
node_resources	Scalar	Number of node resources	required
max_edges	Scalar	Maximum number of edges	required
min_node_resources	Scalar	Minimum number of node resources	required
max_node_resources	Scalar	Maximum number of node resources	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for VONE.

    Args:
        node_resources (chex.Scalar): Number of node resources
        max_edges (chex.Scalar): Maximum number of edges
        min_node_resources (chex.Scalar): Minimum number of node resources
        max_node_resources (chex.Scalar): Maximum number of node resources
    """
    node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_edges: chex.Scalar = struct.field(pytree_node=False)
    min_node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_node_resources: chex.Scalar = struct.field(pytree_node=False)

VONEEnvState

Bases: RSAEnvState

Dataclass to hold environment state for VONE.

Parameters:

Name	Type	Description	Default
node_capacity_array	Array	Node capacity array	required
node_resource_array	Array	Node resource array	required
node_departure_array	Array	Node departure array	required
action_counter	Array	Action counter	required
action_history	Array	Action history	required
node_mask_s	Array	Node mask for source node	required
node_mask_d	Array	Node mask for destination node	required
virtual_topology_patterns	Array	Virtual topology patterns	required
values_nodes	Array	Values for nodes	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvState(RSAEnvState):
    """Dataclass to hold environment state for VONE.

    Args:
        node_capacity_array (chex.Array): Node capacity array
        node_resource_array (chex.Array): Node resource array
        node_departure_array (chex.Array): Node departure array
        action_counter (chex.Array): Action counter
        action_history (chex.Array): Action history
        node_mask_s (chex.Array): Node mask for source node
        node_mask_d (chex.Array): Node mask for destination node
        virtual_topology_patterns (chex.Array): Virtual topology patterns
        values_nodes (chex.Array): Values for nodes
    """
    node_capacity_array: chex.Array
    node_resource_array: chex.Array
    node_departure_array: chex.Array
    action_counter: chex.Array
    action_history: chex.Array
    node_mask_s: chex.Array
    node_mask_d: chex.Array
    virtual_topology_patterns: chex.Array
    values_nodes: chex.Array

aggregate_slots(full_mask, params)

Aggregate slot mask to reduce action space. Only used if the --aggregate_slots flag is set to > 1. Aggregated action is valid if there is one valid slot action within the aggregated action window.

Parameters:

Name	Type	Description	Default
full_mask	Array	slot mask	required
params	EnvParams	environment parameters	required

Returns:

Name	Type	Description
agg_mask	Array	aggregated slot mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def aggregate_slots(full_mask: chex.Array, params: EnvParams) -> chex.Array:
    """Aggregate slot mask to reduce action space. Only used if the --aggregate_slots flag is set to > 1.
    Aggregated action is valid if there is one valid slot action within the aggregated action window.

    Args:
        full_mask: slot mask
        params: environment parameters

    Returns:
        agg_mask: aggregated slot mask
    """

    num_actions = math.ceil(params.link_resources/params.aggregate_slots)
    agg_mask = jnp.zeros((params.k_paths, num_actions), dtype=LARGE_FLOAT_DTYPE)

    def get_max(i, mask_val):
        """Get maximum value of array slice of length aggregate_slots."""
        mask_slice = jax.lax.dynamic_slice(
                mask_val,
                (0, i * params.aggregate_slots,),
                (1,  params.aggregate_slots,),
            )
        max_slice = jnp.max(mask_slice).reshape(1, -1)
        return max_slice

    def update_window_max(i, val):
        """Update ith index 'agg_mask' with max of ith slice of length aggregate_slots from 'full_mask'.

        Args:
            i: increments as += aggregate_slots
            val: tuple of (agg_mask, path_mask, path_index).
        Returns:
            new_agg_mask: agg_mask is updated with max of path_mask for window size aggregate_slots
            mask: mask is unchanged
            path_index: path_index is unchanged
        """
        agg_mask = val[0]
        full_mask = val[1]
        path_index = val[2]
        new_agg_mask = jax.lax.dynamic_update_slice(
            agg_mask,
            get_max(i, full_mask),
            (path_index, i),
        )
        return new_agg_mask, full_mask, path_index

    def apply_to_path_mask(i, val):
        """
        Loop through each path for num_actions steps and get_window_max at each step.

        Args:
            i: path index
            val: tuple of (agg_mask, mask) where mask is original link-slot mask and agg_mask is resulting aggregated mask
        Returns:
            new_agg_mask: agg_mask is updated with aggregated path mask
            mask: mask is unchanged
        """
        val = (
            val[0],  # aggregated mask (to be updated)
            val[1][i].reshape(1, -1),  # mask for path i
            i  # path index
        )
        new_agg_mask = jax.lax.fori_loop(
            0,
            num_actions,
            update_window_max,
            val,
        )[0]
        return new_agg_mask, full_mask

    return jax.lax.fori_loop(
            0,
            params.k_paths,
            apply_to_path_mask,
            (agg_mask, full_mask),
        )

calculate_path_capacity(path_length, min_request=100, scale_factor=1.0, alpha=0.0002, NF=4.5, B=10000000000000.0, R_s=100000000000.0, beta_2=-2.17e-26, gamma=0.0012, L_s=100000.0, lambda0=1.55e-06)

From Nevin JOCN paper 2022: https://discovery.ucl.ac.uk/id/eprint/10175456/1/RL_JOCN_accepted.pdf

Source code in xlron/environments/env_funcs.py

def calculate_path_capacity(
        path_length,
        min_request=100,  # Minimum data rate request size
        scale_factor=1.0,  # Scale factor for link capacity
        alpha=0.2e-3, # Fibre attenuation coefficient
        NF=4.5,  # Amplifier noise figure
        B=10e12,  # Total modulated bandwidth
        R_s=100e9,  # Symbol rate
        beta_2=-21.7e-27,  # Dispersion parameter
        gamma=1.2e-3,  # Nonlinear coefficient
        L_s=100e3,  # Span length
        lambda0=1550e-9,  # Wavelength
):
    """From Nevin JOCN paper 2022: https://discovery.ucl.ac.uk/id/eprint/10175456/1/RL_JOCN_accepted.pdf"""
    alpha_lin = alpha / 4.343  # Linear attenuation coefficient
    N_spans = jnp.floor(path_length * 1e3 / L_s)  # Number of fibre spans on path
    L_eff = (1 - jnp.exp(-alpha_lin * L_s)) / alpha_lin  # Effective length of span in m
    sigma_2_ase = (jnp.exp(alpha_lin * L_s) - 1) * 10**(NF/10) * 6.626e-34 * 2.998e8 * R_s / lambda0  # ASE noise power
    span_NSR = jnp.cbrt(2 * sigma_2_ase**2 * alpha_lin * gamma**2 * L_eff**2 *
                        jnp.log(jnp.pi**2 * jnp.abs(beta_2) * B**2 / alpha_lin) / (jnp.pi * jnp.abs(beta_2) * R_s**2))  # Noise-to-signal ratio per span
    path_NSR = jnp.where(N_spans < 1, 1, N_spans) * span_NSR  # Noise-to-signal ratio per path
    path_capacity = 2 * R_s/1e9 * jnp.log2(1 + 1/path_NSR)  # Capacity of path in Gbps
    # Round link capacity down to nearest increment of minimum request size and apply scale factor
    path_capacity = jnp.floor(path_capacity * scale_factor / min_request) * min_request
    return path_capacity

calculate_path_stats(state, params, request)

For use in DeepRMSA agent observation space. Calculate: 1. Size of 1st suitable free spectrum block 2. Index of 1st suitable free spectrum block 3. Required slots on path 4. Avg. free block size 5. Free slots

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
stats	Array	Array of calculated path statistics

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def calculate_path_stats(state: EnvState, params: EnvParams, request: chex.Array) -> chex.Array:
    """For use in DeepRMSA agent observation space.
    Calculate:
    1. Size of 1st suitable free spectrum block
    2. Index of 1st suitable free spectrum block
    3. Required slots on path
    4. Avg. free block size
    5. Free slots

    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        stats: Array of calculated path statistics
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_val = jnp.zeros((params.k_paths, 5), dtype=LARGE_FLOAT_DTYPE)
    # TODO - check if the normalisation is useful
    def body_fun(i, val):
        link_resources = jnp.array(params.link_resources, dtype=LARGE_FLOAT_DTYPE)
        slot_size = jnp.array(params.slot_size, dtype=LARGE_FLOAT_DTYPE)
        slots = get_path_slots(state.link_slot_array, params, nodes_sd, i)
        se = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else jnp.array([1], dtype=SMALL_INT_DTYPE)
        req_slots = jnp.squeeze(required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband))
        req_slots_norm = req_slots * slot_size / jnp.max(params.values_bw.val)
        free_slots_norm = jnp.sum(jnp.where(slots == zero, one, zero), promote_integers=False) / link_resources
        block_sizes = find_block_sizes(slots)
        first_block_index = jnp.argmax(block_sizes >= req_slots)
        first_block_index_norm = first_block_index.astype(LARGE_FLOAT_DTYPE) / link_resources
        first_block_size_norm = jnp.squeeze(
            jax.lax.dynamic_slice(block_sizes, (first_block_index,), (1,))
        ) / req_slots.astype(LARGE_FLOAT_DTYPE)
        avg_block_size_norm = (jnp.sum(block_sizes) /
                               jnp.max(jnp.array([jnp.sum(find_block_starts(slots), promote_integers=False), 1])) /
                               req_slots)
        val = jax.lax.dynamic_update_slice(
            val,
            jnp.array([[
                first_block_size_norm,
                first_block_index_norm,
                req_slots_norm,
                avg_block_size_norm.astype(LARGE_FLOAT_DTYPE),
                free_slots_norm
            ]]),
            (i, 0)
        )  # N.B. that all values are normalised
        return val

    stats = jax.lax.fori_loop(
            0,
            params.k_paths,
            body_fun,
            init_val,
        )

    return stats

check_action_rmsa_gn_model(state, action, params)

Check if action is valid for RSA GN model Args: state (EnvState): Environment state params (EnvParams): Environment parameters action (chex.Array): Action array Returns: bool: True if action is invalid, False if action is valid

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def check_action_rmsa_gn_model(state: EnvState, action: Optional[chex.Array], params: EnvParams) -> bool:
    """Check if action is valid for RSA GN model
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        action (chex.Array): Action array
    Returns:
        bool: True if action is invalid, False if action is valid
    """
    # Check if action is valid
    # TODO - log failure reasons in info
    snr_sufficient_check = check_snr_sufficient(state, params)
    spectrum_reuse_check = check_no_spectrum_reuse(state.link_slot_array)
    # jax.debug.print("spectrum_reuse_check {}", spectrum_reuse_check, ordered=True)
    # jax.debug.print("snr_sufficient_check {}", snr_sufficient_check, ordered=True)
    return jnp.any(jnp.stack((
        spectrum_reuse_check,
        snr_sufficient_check,
    )))

check_action_rsa(state)

Check if action is valid. Combines checks for: - no spectrum reuse

Parameters:

Name	Type	Description	Default
state		current state	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_action_rsa(state):
    """Check if action is valid.
    Combines checks for:
    - no spectrum reuse

    Args:
        state: current state

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(jnp.stack((
        check_no_spectrum_reuse(state.link_slot_array),
    )))

check_action_rwalr(state, action, params)

Combines checks for: - no spectrum reuse - lightpath available and existing

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_action_rwalr(state: EnvState, action: chex.Array, params: EnvParams) -> bool:
    """Combines checks for:
    - no spectrum reuse
    - lightpath available and existing

    Args:
        state: Environment state

    Returns:
        bool: True if check failed, False if check passed

    """
    return jnp.any(jnp.stack((
        check_no_spectrum_reuse(state.link_slot_array),
        jnp.logical_not(check_lightpath_available_and_existing(state, params, action)[0]),
    )))

check_all_nodes_assigned(node_departure_array, total_requested_nodes)

Count negative values on each node (row) in node departure array, sum them, must equal total requested_nodes.

Parameters:

Name	Type	Description	Default
node_departure_array	Array	Node departure array (N x R) where N is number of nodes and R is number of resources	required
total_requested_nodes	int	Total requested nodes (int)	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_all_nodes_assigned(node_departure_array: chex.Array, total_requested_nodes: int) -> bool:
    """Count negative values on each node (row) in node departure array, sum them, must equal total requested_nodes.

    Args:
        node_departure_array: Node departure array (N x R) where N is number of nodes and R is number of resources
        total_requested_nodes: Total requested nodes (int)

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.sum(jnp.sum(jnp.where(node_departure_array < 0, 1, 0), axis=1)) != total_requested_nodes

check_lightpath_available_and_existing(state, params, action)

Check if lightpath is available and existing. Available means that the lightpath does not use slots occupied by a different lightpath. Existing means that the lightpath has already been established.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Name	Type	Description
lightpath_available_check	Tuple[Array, Array, Array, Array]	True if lightpath is available

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def check_lightpath_available_and_existing(state: EnvState, params: EnvParams, action: chex.Array) -> (
        Tuple)[chex.Array, chex.Array, chex.Array, chex.Array]:
    """Check if lightpath is available and existing.
    Available means that the lightpath does not use slots occupied by a different lightpath.
    Existing means that the lightpath has already been established.

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        lightpath_available_check: True if lightpath is available
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_index, initial_slot_index = process_path_action(state, params, action)
    path = get_paths(params, nodes_sd)[path_index]
    # Get unique lightpath index
    lightpath_index = get_lightpath_index(params, nodes_sd, path_index)
    # Get mask for slots that lightpath will occupy
    # negative numbers used so as not to conflict with lightpath indices
    new_lightpath_mask = vmap_set_path_links(
        jnp.full((params.num_links, 1), -2), path, 0, 1, -1
    )
    path_index_array = state.path_index_array[:, initial_slot_index].reshape(-1, 1)
    masked_path_index_array = jnp.where(
        new_lightpath_mask == -1, path_index_array, -2
    )
    lightpath_mask = jnp.where(
        path_index_array == lightpath_index, -1, -2
    )  # Allow current lightpath
    lightpath_existing_check = jnp.array_equal(lightpath_mask, new_lightpath_mask)  # True if all slots are same
    lightpath_mask = jnp.where(masked_path_index_array == -1, -1, lightpath_mask)  # Allow empty slots
    # True if all slots are same or empty
    lightpath_available_check = jnp.logical_or(
        jnp.array_equal(lightpath_mask, new_lightpath_mask), lightpath_existing_check
    )
    curr_lightpath_capacity = jnp.max(
        jnp.where(new_lightpath_mask == -1, state.link_capacity_array[:, initial_slot_index].reshape(-1, 1), 0)
    )
    return lightpath_available_check, lightpath_existing_check, curr_lightpath_capacity, lightpath_index

check_min_two_nodes_assigned(node_departure_array)

Count negative values on each node (row) in node departure array, sum them, must be 2 or greater. This check is important if e.g. an action contains 2 nodes the same therefore only assigns 1 node. Return False if check passed, True if check failed

Parameters:

Name	Type	Description	Default
node_departure_array	Array	Node departure array (N x R) where N is number of nodes and R is number of resources	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_min_two_nodes_assigned(node_departure_array: chex.Array):
    """Count negative values on each node (row) in node departure array, sum them, must be 2 or greater.
    This check is important if e.g. an action contains 2 nodes the same therefore only assigns 1 node.
    Return False if check passed, True if check failed

    Args:
        node_departure_array: Node departure array (N x R) where N is number of nodes and R is number of resources

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.sum(jnp.sum(jnp.where(node_departure_array < 0, 1, 0), axis=1)) <= 1

check_no_spectrum_reuse(link_slot_array)

slot-=1 when used, should be zero when unoccupied, so check if any < -1 in slot array.

Parameters:

Name	Type	Description	Default
link_slot_array		Link slot array (L x S) where L is number of links and S is number of slots	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_no_spectrum_reuse(link_slot_array):
    """slot-=1 when used, should be zero when unoccupied, so check if any < -1 in slot array.

    Args:
        link_slot_array: Link slot array (L x S) where L is number of links and S is number of slots

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(link_slot_array < -1)

check_node_capacities(capacity_array)

Sum selected nodes array and check less than node resources.

Parameters:

Name	Type	Description	Default
capacity_array	Array	Node capacity array (N x 1) where N is number of nodes	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_node_capacities(capacity_array: chex.Array) -> bool:
    """Sum selected nodes array and check less than node resources.

    Args:
        capacity_array: Node capacity array (N x 1) where N is number of nodes

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(capacity_array < 0)

check_snr_sufficient(state, params)

Check if SNR is sufficient for all connections Args: state (EnvState): Environment state params (EnvParams): Environment parameters Returns: jnp.array: 1 if SNR is sufficient for connection else 0

Source code in xlron/environments/env_funcs.py

def check_snr_sufficient(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> chex.Array:
    """Check if SNR is sufficient for all connections
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: 1 if SNR is sufficient for connection else 0
    """
    # TODO - this check needs to be faster!
    required_snr_array = get_required_snr_se_kurtosis_array(state.modulation_format_index_array, 2, params)
    # Transform lightpath index array by getting lightpath value, getting path-link array, and summing inverse link SNRs
    lightpath_snr_array = get_lightpath_snr(state, params)
    check_snr_sufficient = jnp.where(lightpath_snr_array >= required_snr_array, 0, 1)
    # jax.debug.print("check_snr_sufficient {}", check_snr_sufficient, ordered=True)
    # jax.debug.print("required_snr_array {}", required_snr_array, ordered=True)
    # jax.debug.print("lightpath_snr_array {}", lightpath_snr_array, ordered=True)
    # jax.debug.print("state.modulation_format_index_array {}", state.modulation_format_index_array, ordered=True)
    # jax.debug.print("state.channel_centre_bw_array {}", state.channel_centre_bw_array, ordered=True)
    # jax.debug.print("state.channel_power_array {}", state.channel_power_array, ordered=True)
    return jnp.any(check_snr_sufficient)

check_topology(action_history, topology_pattern)

Check that each unique virtual node (as indicated by topology pattern) is assigned to a consistent physical node i.e. start and end node of ring is same physical node. Method: For each node index in topology pattern, mask action history with that index, then find max value in masked array. If max value is not the same for all values for that virtual node in action history, then return 1, else 0. Array should be all zeroes at the end, so do an any() check on that. e.g. virtual topology pattern = [2,1,3,1,4,1,2] 3 node ring action history = [0,34,4,0,3,1,0] meaning v node "2" goes p node 0, v node "3" goes p node 4, v node "4" goes p node 3 The numbers in-between relate to the slot action. If any value in the array is 1, a virtual node is assigned to multiple different physical nodes. Need to check from both perspectives: 1. For each virtual node, check that all physical nodes are the same 2. For each physical node, check that all virtual nodes are the same

Parameters:

Name	Type	Description	Default
action_history		Action history	required
topology_pattern		Topology pattern	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_topology(action_history, topology_pattern):
    """Check that each unique virtual node (as indicated by topology pattern) is assigned to a consistent physical node
    i.e. start and end node of ring is same physical node.
    Method:
    For each node index in topology pattern, mask action history with that index, then find max value in masked array.
    If max value is not the same for all values for that virtual node in action history, then return 1, else 0.
    Array should be all zeroes at the end, so do an any() check on that.
    e.g. virtual topology pattern = [2,1,3,1,4,1,2]  3 node ring
    action history = [0,34,4,0,3,1,0]
    meaning v node "2" goes p node 0, v node "3" goes p node 4, v node "4" goes p node 3
    The numbers in-between relate to the slot action.
    If any value in the array is 1, a virtual node is assigned to multiple different physical nodes.
    Need to check from both perspectives:
    1. For each virtual node, check that all physical nodes are the same
    2. For each physical node, check that all virtual nodes are the same

    Args:
        action_history: Action history
        topology_pattern: Topology pattern

    Returns:
        bool: True if check failed, False if check passed
    """
    def loop_func_virtual(i, val):
        # Get indices of physical node in action history that correspond to virtual node i
        masked_val = jnp.where(i == topology_pattern, val, -1)
        # Get maximum value at those indices (should all be same)
        max_node = jnp.max(masked_val)
        # For relevant indices, if max value then return 0 else 1
        val = jnp.where(masked_val != -1, masked_val != max_node, val)
        return val
    def loop_func_physical(i, val):
        # Get indices of virtual nodes in topology pattern that correspond to physical node i
        masked_val = jnp.where(i == action_history, val, -1)
        # Get maximum value at those indices (should all be same)
        max_node = jnp.max(masked_val)
        # For relevant indices, if max value then return 0 else 1
        val = jnp.where(masked_val != -1, masked_val != max_node, val)
        return val
    topology_pattern = topology_pattern[::2]  # Only look at node indices, not slot actions
    action_history = action_history[::2]
    check_virtual = jax.lax.fori_loop(jnp.min(topology_pattern), jnp.max(topology_pattern)+1, loop_func_virtual, action_history)
    check_physical = jax.lax.fori_loop(jnp.min(action_history), jnp.max(action_history)+1, loop_func_physical, topology_pattern)
    check = jnp.concatenate((check_virtual, check_physical))
    return jnp.any(check)

check_unique_nodes(node_departure_array)

Count negative values on each node (row) in node departure array, must not exceed 1.

Parameters:

Name	Type	Description	Default
node_departure_array	Array	Node departure array (N x R) where N is number of nodes and R is number of resources	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

@jax.jit
def check_unique_nodes(node_departure_array: chex.Array) -> bool:
    """Count negative values on each node (row) in node departure array, must not exceed 1.

    Args:
        node_departure_array: Node departure array (N x R) where N is number of nodes and R is number of resources

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(jnp.sum(jnp.where(node_departure_array < zero, one, zero), axis=1, promote_integers=False) > one)

check_vone_action(state, remaining_actions, total_requested_nodes)

Check if action is valid. Combines checks for: - sufficient node capacities - unique nodes assigned - minimum two nodes assigned - all requested nodes assigned - correct topology pattern - no spectrum reuse

Parameters:

Name	Type	Description	Default
state		current state	required
remaining_actions		remaining actions	required
total_requested_nodes		total requested nodes	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_vone_action(state, remaining_actions, total_requested_nodes):
    """Check if action is valid.
    Combines checks for:
    - sufficient node capacities
    - unique nodes assigned
    - minimum two nodes assigned
    - all requested nodes assigned
    - correct topology pattern
    - no spectrum reuse

    Args:
        state: current state
        remaining_actions: remaining actions
        total_requested_nodes: total requested nodes

    Returns:
        bool: True if check failed, False if check passed
    """
    checks = jnp.stack((
        check_node_capacities(state.node_capacity_array),
        check_unique_nodes(state.node_departure_array),
        # TODO (VONE) - Remove two nodes check if impairs performance
        #  (check_all_nodes_assigned is sufficient but fails after last action of request instead of earlier)
        check_min_two_nodes_assigned(state.node_departure_array),
        jax.lax.cond(
            jnp.equal(remaining_actions, jnp.array(1)),
            lambda x: check_all_nodes_assigned(*x),
            lambda x: jnp.array(False),
            (state.node_departure_array, total_requested_nodes)
        ),
        jax.lax.cond(
            jnp.equal(remaining_actions, jnp.array(1)),
            lambda x: check_topology(*x),
            lambda x: jnp.array(False),
            (state.action_history, state.request_array[1])
        ),
        check_no_spectrum_reuse(state.link_slot_array),
    ))
    #jax.debug.print("Checks: {}", checks, ordered=True)
    return jnp.any(checks)

convert_node_probs_to_traffic_matrix(node_probs)

Convert list of node probabilities to symmetric traffic matrix.

Parameters:

Name	Type	Description	Default
node_probs	list	node probabilities	required

Returns:

Name	Type	Description
traffic_matrix	Array	traffic matrix

Source code in xlron/environments/env_funcs.py

def convert_node_probs_to_traffic_matrix(node_probs: list) -> chex.Array:
    """Convert list of node probabilities to symmetric traffic matrix.

    Args:
        node_probs: node probabilities

    Returns:
        traffic_matrix: traffic matrix
    """
    matrix = jnp.outer(node_probs, node_probs).astype(SMALL_FLOAT_DTYPE)
    # Set lead diagonal to zero
    matrix = jnp.where(jnp.eye(matrix.shape[0]) == 1, 0, matrix)
    matrix = normalise_traffic_matrix(matrix)
    return matrix

create_run_name(config)

Create name for run based on config flags

Source code in xlron/environments/env_funcs.py

def create_run_name(config: Union[box.Box, dict]) -> str:
    """Create name for run based on config flags"""
    env_type = config["env_type"]
    topology = config["topology_name"]
    slots = config["link_resources"]
    gnn = "_GNN" if config["USE_GNN"] else ""
    incremental = "_INC" if config["incremental_loading"] else ""
    run_name = f"{env_type}_{topology}_{slots}{gnn}{incremental}".upper()
    if config["EVAL_HEURISTIC"]:
        run_name += f"_{config['path_heuristic']}"
        if env_type.lower() == "vone":
            run_name += f"_{config['node_heuristic']}"
    elif config["EVAL_MODEL"]:
        run_name += f"_EVAL"
    return run_name

decrement_action_counter(state)

Decrement action counter in-place. Used in VONE environments.

Source code in xlron/environments/env_funcs.py

def decrement_action_counter(state):
    """Decrement action counter in-place. Used in VONE environments."""
    state.action_counter.at[-1].add(-1)
    return state

finalise_action_rsa(state, params)

Turn departure times positive.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def finalise_action_rsa(state: EnvState, params: Optional[EnvParams]):
    """Turn departure times positive.

    Args:
        state: current state

    Returns:
        state: updated state
    """
    _, requested_datarate = read_rsa_request(state.request_array)
    state = state.replace(
        link_slot_departure_array=make_positive(state.link_slot_departure_array),
        accepted_services=state.accepted_services + 1,
        accepted_bitrate=state.accepted_bitrate + requested_datarate[0],
        total_bitrate=state.total_bitrate + requested_datarate[0]
    )
    return state

finalise_action_rwalr(state, params)

Turn departure times positive.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def finalise_action_rwalr(state: EnvState, params: Optional[EnvParams]):
    """Turn departure times positive.

    Args:
        state: current state

    Returns:
        state: updated state
    """
    _, requested_datarate = read_rsa_request(state.request_array)
    state = state.replace(
        link_slot_departure_array=make_positive(state.link_slot_departure_array),
        accepted_services=state.accepted_services + 1,
        accepted_bitrate=state.accepted_bitrate + requested_datarate[0],
        total_bitrate=state.total_bitrate + requested_datarate[0]
    )
    return state

finalise_vone_action(state)

Turn departure times positive.

Parameters:

Name	Type	Description	Default
state		current state	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def finalise_vone_action(state):
    """Turn departure times positive.

    Args:
        state: current state

    Returns:
        state: updated state
    """
    state = state.replace(
        node_departure_array=make_positive(state.node_departure_array),
        link_slot_departure_array=make_positive(state.link_slot_departure_array),
        accepted_services=state.accepted_services + 1,
        accepted_bitrate=state.accepted_bitrate  # TODO - get sum of bitrates for requested links
    )
    return state

format_vone_slot_request(state, action)

Format slot request for VONE action into format (source-node, slot, destination-node).

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required
action	Array	action to format	required

Returns:

Type	Description
Array	chex.Array: formatted request

Source code in xlron/environments/env_funcs.py

def format_vone_slot_request(state: EnvState, action: chex.Array) -> chex.Array:
    """Format slot request for VONE action into format (source-node, slot, destination-node).

    Args:
        state: current state
        action: action to format

    Returns:
        chex.Array: formatted request
    """
    remaining_actions = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.action_counter, 2, 1))
    full_request = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.request_array, 0, 1))
    unformatted_request = jax.lax.dynamic_slice_in_dim(full_request, (remaining_actions - 1) * 2, 3)
    node_s = jax.lax.dynamic_slice_in_dim(action, 0, 1)
    requested_slots = jax.lax.dynamic_slice_in_dim(unformatted_request, 1, 1)
    node_d = jax.lax.dynamic_slice_in_dim(action, 2, 1)
    formatted_request = jnp.concatenate((node_s, requested_slots, node_d))
    return formatted_request

generate_arrival_holding_times(key, params)

Generate arrival and holding times based on Poisson distributed events. To understand how sampling from e^-x can be transformed to sample from lambdae^-(x/lambda) see: https://en.wikipedia.org/wiki/Inverse_transform_sampling#Examples Basically, inverse transform sampling is used to sample from a distribution with CDF F(x). The CDF of the exponential distribution (lambdae^-{lambdax}) is F(x) = 1 - e^-{lambdax}. Therefore, the inverse CDF is x = -ln(1-u)/lambda, where u is sample from uniform distribution. Therefore, we need to divide jax.random.exponential() by lambda in order to scale the standard exponential CDF. Experimental histograms of this method compared to random.expovariate() in Python's random library show that the two methods are equivalent. Also see: https://numpy.org/doc/stable/reference/random/generated/numpy.random.exponential.html https://jax.readthedocs.io/en/latest/_autosummary/jax.random.exponential.html

Parameters:

Name	Type	Description	Default
key		PRNG key	required
params		Environment parameters	required

Returns:

Name	Type	Description
arrival_time		Arrival time
holding_time		Holding time

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def generate_arrival_holding_times(key, params):
    """
    Generate arrival and holding times based on Poisson distributed events.
    To understand how sampling from e^-x can be transformed to sample from lambda*e^-(x/lambda) see:
    https://en.wikipedia.org/wiki/Inverse_transform_sampling#Examples
    Basically, inverse transform sampling is used to sample from a distribution with CDF F(x).
    The CDF of the exponential distribution (lambda*e^-{lambda*x}) is F(x) = 1 - e^-{lambda*x}.
    Therefore, the inverse CDF is x = -ln(1-u)/lambda, where u is sample from uniform distribution.
    Therefore, we need to divide jax.random.exponential() by lambda in order to scale the standard exponential CDF.
    Experimental histograms of this method compared to random.expovariate() in Python's random library show that
    the two methods are equivalent.
    Also see: https://numpy.org/doc/stable/reference/random/generated/numpy.random.exponential.html
    https://jax.readthedocs.io/en/latest/_autosummary/jax.random.exponential.html

    Args:
        key: PRNG key
        params: Environment parameters

    Returns:
        arrival_time: Arrival time
        holding_time: Holding time
    """
    key_arrival, key_holding = jax.random.split(key, 2)
    arrival_time = jax.random.exponential(key_arrival, shape=(1,), dtype=SMALL_FLOAT_DTYPE) \
                   / params.arrival_rate  # Divide because it is rate (lambda)
    if params.truncate_holding_time:
        # For DeepRMSA, need to generate holding times that are less than 2*mean_service_holding_time
        key_holding = jax.random.split(key, 5)
        holding_times = jax.vmap(lambda x: jax.random.exponential(x, shape=(1,)) \
                                * params.mean_service_holding_time)(key_holding)
        holding_times = jnp.where(holding_times < 2*params.mean_service_holding_time, holding_times, zero)
        # Get first non-zero value in holding_times
        non_zero_index = jnp.nonzero(holding_times, size=1)[0][0]
        holding_time = jax.lax.dynamic_slice(jnp.squeeze(holding_times), (non_zero_index,), (1,))
    else:
        holding_time = jax.random.exponential(key_holding, shape=(1,), dtype=SMALL_FLOAT_DTYPE) \
                       * params.mean_service_holding_time  # Multiply because it is mean (1/lambda)
    return arrival_time, holding_time

generate_vone_request(key, state, params)

Generate a new request for the VONE environment. The request has two rows. The first row shows the node and slot values. The first three elements of the second row show the number of unique nodes, the total number of steps, and the remaining steps. These first three elements comprise the action counter. The remaining elements of the second row show the virtual topology pattern, i.e. the connectivity of the virtual topology.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def generate_vone_request(key: chex.PRNGKey, state: EnvState, params: EnvParams):
    """Generate a new request for the VONE environment.
    The request has two rows. The first row shows the node and slot values.
    The first three elements of the second row show the number of unique nodes, the total number of steps, and the remaining steps.
    These first three elements comprise the action counter.
    The remaining elements of the second row show the virtual topology pattern, i.e. the connectivity of the virtual topology.
    """
    shape = params.max_edges*2+1  # shape of request array
    key_topology, key_node, key_slot, key_times = jax.random.split(key, 4)
    # Randomly select topology, node resources, slot resources
    pattern = jax.random.choice(key_topology, state.virtual_topology_patterns)
    action_counter = jax.lax.dynamic_slice(pattern, (0,), (3,))
    topology_pattern = jax.lax.dynamic_slice(pattern, (3,), (pattern.shape[0]-3,))
    selected_node_values = jax.random.choice(key_node, state.values_nodes, shape=(shape,))
    selected_bw_values = jax.random.choice(key_slot, params.values_bw.val, shape=(shape,))
    # Create a mask for odd and even indices
    mask = jnp.tile(jnp.array([0, 1]), (shape+1) // 2)[:shape]
    # Vectorized conditional replacement using mask
    first_row = jnp.where(mask, selected_bw_values, selected_node_values)
    # Make sure node request values are consistent for same virtual nodes
    first_row = jax.lax.fori_loop(
        2,  # Lowest node index in virtual topology requests is 2
        shape,  # Highest possible node index in virtual topology requests is shape-1
        lambda i, x: jnp.where(topology_pattern == i, selected_node_values[i], x),
        first_row
    )
    # Mask out unused part of request array
    first_row = jnp.where(topology_pattern == 0, 0, first_row)
    # Set times
    arrival_time, holding_time = generate_arrival_holding_times(key, params)
    state = state.replace(
        holding_time=holding_time,
        current_time=state.current_time + arrival_time,
        action_counter=action_counter,
        request_array=jnp.vstack((first_row, topology_pattern)),
        action_history=init_action_history(params),
        total_requests=state.total_requests + 1
    )
    state = remove_expired_node_requests(state, params) if not params.incremental_loading else state
    state = remove_expired_services_rsa(state, params) if not params.incremental_loading else state
    return state

get_best_modulation_format(state, path, initial_slot_index, launch_power, params)

Get best modulation format for lightpath. "Best" is the highest order that has SNR requirements below available. Try each modulation format, calculate SNR for each, then return the highest order possible. Args: state (EnvState): Environment state path (chex.Array): Path array initial_slot_index (int): Initial slot index params (EnvParams): Environment parameters Returns: jnp.array: Acceptable modulation format indices

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(3,))
def get_best_modulation_format(state: EnvState, path: chex.Array, initial_slot_index: int, launch_power: chex.Array, params: EnvParams) -> chex.Array:
    """Get best modulation format for lightpath. "Best" is the highest order that has SNR requirements below available.
    Try each modulation format, calculate SNR for each, then return the highest order possible.
    Args:
        state (EnvState): Environment state
        path (chex.Array): Path array
        initial_slot_index (int): Initial slot index
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Acceptable modulation format indices
    """
    _, requested_datarate = read_rsa_request(state.request_array)
    mod_format_count = params.modulations_array.val.shape[0]
    acceptable_mod_format_indices = jnp.full((mod_format_count,), -2)

    def acceptable_modulation_format(i, acceptable_format_indices):
        req_snr = params.modulations_array.val[i][2] + params.snr_margin
        se = params.modulations_array.val[i][1]
        req_slots = required_slots(requested_datarate, se, params.slot_size, params.guardband)
        # TODO - need to check we don't overwrite values in already occupied slots
        # Possible approaches:
        # Check slot occupancy? Probably would need to iterate through for num_slots, but that's an issue
        # What about we allocate and then fix up later, e.g. could it be possible to just add the modulation format on top without
        # check sum of path links prior to assigning?
        #
        new_state = state.replace(
            channel_power_array=vmap_set_path_links(
                state.channel_power_array, path, initial_slot_index, req_slots, launch_power),
            channel_centre_bw_array=vmap_set_path_links(
                state.channel_centre_bw_array, path, initial_slot_index, req_slots, params.slot_size)
        )
        snr_value = get_minimum_snr_of_channels_on_path(new_state, path, initial_slot_index, req_slots, params)
        # jax.debug.print("snr_value {}", snr_value, ordered=True)
        # jax.debug.print("req_snr {}", req_snr, ordered=True)
        acceptable_format_index = jnp.where(snr_value >= req_snr, i, -1).reshape((1,))
        acceptable_format_indices = jax.lax.dynamic_update_slice(acceptable_format_indices, acceptable_format_index, (i,))
        # jax.debug.print("acceptable_format_indices {}", acceptable_format_indices, ordered=True)
        return acceptable_format_indices

    acceptable_mod_format_indices = jax.lax.fori_loop(
        0,
        mod_format_count,
        acceptable_modulation_format,
        acceptable_mod_format_indices
    )
    return acceptable_mod_format_indices

get_best_modulation_format_simple(state, path, initial_slot_index, params)

Get modulation format for lightpath. Assume worst case (least Gaussian) modulation format when calculating SNR. Args: state (EnvState): Environment state path (chex.Array): Path array initial_slot_index (int): Initial slot index params (EnvParams): Environment parameters Returns: jnp.array: Acceptable modulation format indices

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(3,))
def get_best_modulation_format_simple(
        state: RSAGNModelEnvState, path: chex.Array, initial_slot_index: int, params: RSAGNModelEnvParams
) -> chex.Array:
    """Get modulation format for lightpath.
    Assume worst case (least Gaussian) modulation format when calculating SNR.
    Args:
        state (EnvState): Environment state
        path (chex.Array): Path array
        initial_slot_index (int): Initial slot index
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Acceptable modulation format indices
    """
    link_snr_array = get_snr_link_array(state, params)
    snr_value = get_snr_for_path(path, link_snr_array, params)[initial_slot_index] - params.snr_margin  # Margin
    mod_format_count = params.modulations_array.val.shape[0]
    acceptable_mod_format_indices = jnp.arange(mod_format_count)
    req_snr = params.modulations_array.val[:, 2] + params.snr_margin
    acceptable_mod_format_indices = jnp.where(snr_value >= req_snr,
                                              acceptable_mod_format_indices,
                                              jnp.full((mod_format_count,), -2))
    return acceptable_mod_format_indices

get_centre_frequency(initial_slot_index, num_slots, params)

Get centre frequency for new lightpath

Parameters:

Name	Type	Description	Default
initial_slot_index	Array	Centre frequency of first slot	required
num_slots	float	Number of slots	required
params	RSAGNModelEnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Centre frequency for new lightpath

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def get_centre_frequency(initial_slot_index: int, num_slots: int, params: RSAGNModelEnvParams) -> chex.Array:
    """Get centre frequency for new lightpath

    Args:
        initial_slot_index (chex.Array): Centre frequency of first slot
        num_slots (float): Number of slots
        params (RSAGNModelEnvParams): Environment parameters

    Returns:
        chex.Array: Centre frequency for new lightpath
    """
    slot_centres = (jnp.arange(params.link_resources) - (params.link_resources - 1) / 2) * params.slot_size
    return slot_centres[initial_slot_index] + ((params.slot_size * (num_slots - 1)) / 2)

get_edge_disjoint_paths(graph)

Get edge disjoint paths between all nodes in graph.

Parameters:

Name	Type	Description	Default
graph	Graph	graph	required

Returns:

Name	Type	Description
dict	dict	edge disjoint paths (path is list of edges)

Source code in xlron/environments/env_funcs.py

def get_edge_disjoint_paths(graph: nx.Graph) -> dict:
    """Get edge disjoint paths between all nodes in graph.

    Args:
        graph: graph

    Returns:
        dict: edge disjoint paths (path is list of edges)
    """
    result = {n: {} for n in graph}
    for n1, n2 in itertools.combinations(graph, 2):
        # Sort by number of links in path
        # TODO - sort by path length
        result[n1][n2] = sorted(list(nx.edge_disjoint_paths(graph, n1, n2)), key=len)
        result[n2][n1] = sorted(list(nx.edge_disjoint_paths(graph, n2, n1)), key=len)
    return result

get_launch_power(state, path_action, power_action, params)

Get launch power for new lightpath. N.B. launch power is specified in dBm but is converted to linear units when stored in channel_power_array. This func returns linear units (mW). Path action is used to determine the launch power in the case of tabular launch power type. Power action is used to determine the launch power in the case of RL launch power type. During masking, power action is set as state.launch_power_array[0], which is set by the RL agent. Args: state (EnvState): Environment state path_action (chex.Array): Action specifying path index (0 to k_paths-1) power_action (chex.Array): Action specifying launch power in dBm params (EnvParams): Environment parameters Returns: chex.Array: Launch power for new lightpath

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_launch_power(state: EnvState, path_action: chex.Array, power_action: chex.Array, params: EnvParams) -> chex.Array:
    """Get launch power for new lightpath. N.B. launch power is specified in dBm but is converted to linear units
    when stored in channel_power_array. This func returns linear units (mW).
    Path action is used to determine the launch power in the case of tabular launch power type.
    Power action is used to determine the launch power in the case of RL launch power type. During masking,
    power action is set as state.launch_power_array[0], which is set by the RL agent.
    Args:
        state (EnvState): Environment state
        path_action (chex.Array): Action specifying path index (0 to k_paths-1)
        power_action (chex.Array): Action specifying launch power in dBm
        params (EnvParams): Environment parameters
    Returns:
        chex.Array: Launch power for new lightpath
    """
    k_path_index, _ = process_path_action(state, params, path_action)
    if params.launch_power_type == 1:  # Fixed
        return state.launch_power_array[0]
    elif params.launch_power_type == 2:  # Tabular (one row per path)
        nodes_sd, requested_datarate = read_rsa_request(state.request_array)
        source, dest = nodes_sd
        i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
        return state.launch_power_array[i+k_path_index]
    elif params.launch_power_type == 3:  # RL
        return power_action
    elif params.launch_power_type == 4:  # Scaled
        nodes_sd, requested_datarate = read_rsa_request(state.request_array)
        source, dest = nodes_sd
        i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
        # Get path length
        link_length_array = jnp.sum(params.link_length_array.val, axis=1, promote_integers=False)
        path_length = jnp.sum(link_length_array[i+k_path_index], promote_integers=False)
        path_link_array = jnp.unpackbits(params.path_link_array.val)[:, params.num_links] if params.pack_path_bits \
            else params.path_link_array.val
        maximum_path_length = jnp.max(jnp.dot(path_link_array, params.link_length_array.val))
        return state.launch_power_array[0] * (path_length / maximum_path_length)
    else:
        raise ValueError("Invalid launch power type. Check params.launch_power_type")

get_lightpath_snr(state, params)

Get SNR for each link on path. N.B. that in most cases it is more efficient to calculate the SNR for every possible path, rather than a slot-by-slot basis. But in some cases slot-by-slot is better i.e. when kN(N-1)/2 > LS Args: state (RSAGNModelEnvState): Environment state params (RSAGNModelEnvParams): Environment parameters

Returns:

Type	Description
Array	chex.array: SNR for each link on path

Source code in xlron/environments/env_funcs.py

def get_lightpath_snr(state: RSAGNModelEnvParams, params: RSAGNModelEnvParams) -> chex.Array:
    """Get SNR for each link on path.
    N.B. that in most cases it is more efficient to calculate the SNR for every possible path, rather than a slot-by-slot basis.
    But in some cases slot-by-slot is better i.e. when k*N(N-1)/2 > L*S
    Args:
        state (RSAGNModelEnvState): Environment state
        params (RSAGNModelEnvParams): Environment parameters

    Returns:
        chex.array: SNR for each link on path
    """
    # Get the SNR for the channel that the path occupies
    path_snr_array = jax.vmap(get_snr_for_path, in_axes=(0, None, None))(params.path_link_array.val, state.link_snr_array, params)
    # Where value in path_index_array matches index of path_snr_array, substitute in SNR value
    slot_indices = jnp.arange(params.link_resources)
    lightpath_snr_array = jax.vmap(jax.vmap(lambda x, si: path_snr_array[x][si], in_axes=(0, 0)), in_axes=(0, None))(state.path_index_array, slot_indices)
    return lightpath_snr_array

get_minimum_snr_of_channels_on_path(state, path, slot_index, req_slots, params)

Get the minimum value of the SNR on newly assigned channels. N.B. this requires the link_snr_array to have already been calculated and present in state.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def get_minimum_snr_of_channels_on_path(
        state: RSAGNModelEnvState, path: chex.Array, slot_index: chex.Array, req_slots: int, params: RSAGNModelEnvParams
) -> chex.Array:
    """Get the minimum value of the SNR on newly assigned channels.
    N.B. this requires the link_snr_array to have already been calculated and present in state."""
    snr_value_all_channels = get_snr_for_path(path, state.link_snr_array, params)
    min_snr_value_sub_channels = jnp.min(
        jnp.concatenate([
            snr_value_all_channels[slot_index].reshape((1,)),
            snr_value_all_channels[slot_index + req_slots - 1].reshape((1,))
        ], axis=0)
    )
    return min_snr_value_sub_channels

get_num_spectral_features(n_nodes)

Heuristic for number of spectral features based on graph size.

Parameters:

Name	Type	Description	Default
n_nodes	int	Number of nodes in the graph	required

Returns:

Type	Description
int	Number of spectral features to use, clamped between 3 and 15.
int	Follows log2(n_nodes) scaling as reasonable default.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_num_spectral_features(n_nodes: int) -> int:
    """Heuristic for number of spectral features based on graph size.

    Args:
        n_nodes: Number of nodes in the graph

    Returns:
        Number of spectral features to use, clamped between 3 and 15.
        Follows log2(n_nodes) scaling as reasonable default.
    """
    return jnp.minimum(jnp.maximum(3, jnp.floor(jnp.log2(n_nodes))), 15).astype(int)

get_path_from_path_index_array(path_index_array, path_link_array)

Get path from path index array. Args: path_index_array (chex.Array): Path index array path_link_array (chex.Array): Path link array

Returns:

Type	Description
Array	jnp.array: path index values replaced with binary path-link arrays

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_path_from_path_index_array(path_index_array: chex.Array, path_link_array: chex.Array) -> chex.Array:
    """Get path from path index array.
    Args:
        path_index_array (chex.Array): Path index array
        path_link_array (chex.Array): Path link array

    Returns:
        jnp.array: path index values replaced with binary path-link arrays
    """
    # TODO - support unpacking bits (if this function ends up being used)
    def get_index_from_link(link):
        return jax.vmap(lambda x: path_link_array[x], in_axes=(0,))(link)

    return jax.vmap(get_index_from_link, in_axes=(0,))(path_index_array)

get_path_index_array(params, nodes)

Indices of paths between source and destination from path array

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_path_index_array(params, nodes):
    """Indices of paths between source and destination from path array"""
    # get source and destination nodes in order (for accurate indexing of path-link array)
    source, dest = nodes.astype(LARGE_INT_DTYPE)
    i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
    index_array = jax.lax.dynamic_slice(jnp.arange(0, params.path_link_array.shape[0], dtype=LARGE_INT_DTYPE), (i,), (params.k_paths,))
    return index_array

get_path_indices(s, d, k, N, directed=False)

Get path indices for a given source, destination and number of paths. If source > destination and the graph is directed (two fibres per link, one in each direction) then an offset is added to the index to get the path in the other direction (the offset is the total number source-dest pairs).

Parameters:

Name	Type	Description	Default
s	int	Source node index	required
d	int	Destination node index	required
k	int	Number of paths	required
N	int	Number of nodes	required
directed	bool	Whether graph is directed. Defaults to False.	False

Returns:

Type	Description
Array	jnp.array: Start index on path-link array for candidate paths

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2, 3, 4))
def get_path_indices(s: int, d: int, k: int, N: int, directed: bool = False) -> chex.Array:
    """Get path indices for a given source, destination and number of paths.
    If source > destination and the graph is directed (two fibres per link, one in each direction) then an offset is
    added to the index to get the path in the other direction (the offset is the total number source-dest pairs).

    Args:
        s (int): Source node index
        d (int): Destination node index
        k (int): Number of paths
        N (int): Number of nodes
        directed (bool, optional): Whether graph is directed. Defaults to False.

    Returns:
        jnp.array: Start index on path-link array for candidate paths
    """
    node_indices = jnp.arange(N, dtype=LARGE_INT_DTYPE)
    indices_to_s = jnp.where(node_indices < s, node_indices, jnp.array(0, dtype=LARGE_INT_DTYPE))
    indices_to_d = jnp.where(node_indices < d, node_indices, jnp.array(0, dtype=LARGE_INT_DTYPE))
    # If two fibres per link, add offset to index to get fibre in other direction if source > destination
    directed_offset = directed * (s > d) * N * (N - 1) * k / 2
    # The following equation is based on the combinations formula
    forward = ((N * s + d - jnp.sum(indices_to_s, promote_integers=False) - 2 * s - 1) * k)
    backward = ((N * d + s - jnp.sum(indices_to_d, promote_integers=False) - 2 * d - 1) * k)
    return forward * (s < d) + backward * (s > d) + directed_offset.astype(LARGE_INT_DTYPE)

get_path_slots(link_slot_array, params, nodes_sd, i, agg_func='max')

Get slots on each constitutent link of path from link_slot_array (L x S), then aggregate to get (S x 1) representation of slots on path.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	link-slot array	required
params	EnvParams	environment parameters	required
nodes_sd	Array	source-destination nodes	required
i	int	path index	required
agg_func	str	aggregation function (max or sum). If max, result will be available slots on path. If sum, result will contain information on edge features.	'max'

Returns:

Name	Type	Description
slots	Array	slots on path

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 4))
def get_path_slots(link_slot_array: chex.Array, params: EnvParams, nodes_sd: chex.Array, i: int, agg_func: str = "max") -> chex.Array:
    """Get slots on each constitutent link of path from link_slot_array (L x S),
    then aggregate to get (S x 1) representation of slots on path.

    Args:
        link_slot_array: link-slot array
        params: environment parameters
        nodes_sd: source-destination nodes
        i: path index
        agg_func: aggregation function (max or sum).
            If max, result will be available slots on path.
            If sum, result will contain information on edge features.

    Returns:
        slots: slots on path
    """
    path = get_paths(params, nodes_sd)[i]
    path = path.reshape((params.num_links, 1))
    # Get links and collapse to single dimension
    num_slots = params.link_resources if agg_func == "max" else math.ceil(params.link_resources/params.aggregate_slots)
    slots = jnp.where(path, link_slot_array, jnp.zeros(num_slots, dtype=LARGE_FLOAT_DTYPE))
    # Make any -1s positive then get max for each slot across links
    if agg_func == "max":
        # Use this for getting slots from link_slot_array
        slots = jnp.max(jnp.absolute(slots), axis=0)
    elif agg_func == "sum":
        # TODO - consider using an RNN (or S5) to aggregate edge features
        # Use this (or alternative) for aggregating edge features from GNN
        slots = jnp.sum(slots, axis=0, promote_integers=False)
    elif agg_func == "mean":
        # Use this for getting mean value in slot index along path
        slots = jnp.mean(slots, axis=0)
    else:
        raise ValueError("agg_func must be 'max' or 'sum' or 'mean'")
    return slots

get_paths(params, nodes)

Get k paths between source and destination

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_paths(params, nodes):
    """Get k paths between source and destination"""
    index_array = get_path_index_array(params, nodes)
    paths = jnp.take(params.path_link_array.val, index_array, axis=0)
    if params.pack_path_bits:  # Unpack the bit-packed paths
        paths = jnp.unpackbits(paths, axis=1)[:, :params.num_links]
    return paths

get_paths_obs_gn_model(state, params)

Get observation space for launch power optimization (with numerical stability).

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_paths_obs_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> chex.Array:
    # TODO - make this just show the stats from just one path at a time
    """Get observation space for launch power optimization (with numerical stability)."""
    request_array = state.request_array.reshape((-1,))
    path_stats = calculate_path_stats(state, params, request_array)
    # Remove first 3 items of path stats for each path
    path_stats = path_stats[:, 3:]
    link_length_array = jnp.sum(params.link_length_array.val, axis=1, promote_integers=False)
    lightpath_snr_array = get_lightpath_snr(state, params)
    nodes_sd, requested_datarate = read_rsa_request(request_array)
    source, dest = nodes_sd

    def calculate_gn_path_stats(k_path_index, init_val):
        # Get path index
        path_index = get_path_indices(source, dest, params.k_paths, params.num_nodes,
                                      directed=params.directed_graph) + k_path_index
        path_link_array = jnp.unpackbits(params.path_link_array.val, axis=1)[:, :params.num_links] if params.pack_path_bits \
            else params.path_link_array.val
        path = path_link_array[path_index]
        path_length = jnp.dot(path, link_length_array)
        max_path_length = jnp.max(jnp.dot(path_link_array, link_length_array))
        path_length_norm = path_length / max_path_length
        max_path_length_hops = jnp.max(jnp.sum(path_link_array, axis=1, promote_integers=False))
        path_length_hops_norm = jnp.sum(path, promote_integers=False).astype(LARGE_FLOAT_DTYPE) / max_path_length_hops
        # Connections on path
        num_connections = jnp.where(path == 1, jnp.where(state.channel_power_array > 0, one, zero).sum(axis=1), zero).sum()
        num_connections_norm = num_connections / jnp.array(params.link_resources, dtype=LARGE_FLOAT_DTYPE)
        # Mean power of connections on path
        # make path with row length equal to link_resource (+1 to avoid zero division)
        mean_power_norm = (jnp.where(path == one, state.channel_power_array.sum(axis=1), zero).sum() /
                           (jnp.where(num_connections > zero, num_connections, one) * params.max_power))
        # Mean SNR of connections on the path links
        max_snr = jnp.array(50, dtype=LARGE_FLOAT_DTYPE)  # Nominal value for max GSNR in dB
        mean_snr_norm = (jnp.where(path == one, lightpath_snr_array.sum(axis=1), zero).sum(promote_integers=False) /
                         (jnp.where(num_connections > zero, num_connections, one) * max_snr))
        return jax.lax.dynamic_update_slice(
            init_val,
            jnp.array([[
                path_length,
                path_length_hops_norm,
                num_connections_norm,
                mean_power_norm,
                mean_snr_norm
            ]]),
            (k_path_index, 0),
        )

    gn_path_stats = jnp.zeros((params.k_paths, 5), dtype=LARGE_FLOAT_DTYPE)
    gn_path_stats = jax.lax.fori_loop(
        0, params.k_paths, calculate_gn_path_stats, gn_path_stats
    )
    all_stats = jnp.concatenate([path_stats, gn_path_stats], axis=1)
    return jnp.concatenate(
        (
            jnp.array([source]),
            requested_datarate / 100.,
            jnp.array([dest]),
            jnp.reshape(state.holding_time, (-1,)),
            jnp.reshape(all_stats, (-1,)),
        ),
        axis=0,
    )

get_paths_se(params, nodes)

Get max. spectral efficiency of modulation format on k paths between source and destination

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_paths_se(params, nodes):
    """Get max. spectral efficiency of modulation format on k paths between source and destination"""
    # get source and destination nodes in order (for accurate indexing of path-link array)
    index_array = get_path_index_array(params, nodes)
    return jnp.take(params.path_se_array.val, index_array, axis=0)

get_required_snr_se_kurtosis_array(modulation_format_index_array, col_index, params)

Convert modulation format index to required SNR or spectral efficiency. Modulation format index array contains the index of the modulation format used by the channel. The modulation index references a row in the modulations array, which contains SNR and SE values.

Parameters:

Name	Type	Description	Default
modulation_format_index_array	Array	Modulation format index array	required
col_index	int	Column index for required SNR or spectral efficiency	required
params	RSAGNModelEnvParams	Environment parameters	required

Returns:

Type	Description
Array	jnp.array: Required SNR for each channel (min. SNR for empty channel (mod. index 0))

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 2,))
def get_required_snr_se_kurtosis_array(modulation_format_index_array: chex.Array, col_index: int, params: RSAGNModelEnvParams) -> chex.Array:
    """Convert modulation format index to required SNR or spectral efficiency.
    Modulation format index array contains the index of the modulation format used by the channel.
    The modulation index references a row in the modulations array, which contains SNR and SE values.

    Args:
        modulation_format_index_array (chex.Array): Modulation format index array
        col_index (int): Column index for required SNR or spectral efficiency
        params (RSAGNModelEnvParams): Environment parameters

    Returns:
        jnp.array: Required SNR for each channel (min. SNR for empty channel (mod. index 0))
    """
    return jax.vmap(get_required_snr_se_kurtosis_on_link, in_axes=(0, None, None))(modulation_format_index_array, col_index, params)

get_snr_link_array(state, params)

Get SNR per link Args: state (EnvState): Environment state params (EnvParams): Environment parameters Returns: jnp.array: SNR per link

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_snr_link_array(state: EnvState, params: EnvParams) -> chex.Array:
    """Get SNR per link
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: SNR per link
    """

    def get_link_snr(link_index, state, params):
        # Get channel power, channel centre, bandwidth, and noise figure
        link_lengths = params.link_length_array[link_index, :]
        num_spans = jnp.ceil(jnp.sum(link_lengths)*1e3 / params.max_span_length).astype(MED_INT_DTYPE)
        if params.mod_format_correction:
            mod_format_link = state.modulation_format_index_array[link_index, :]
            kurtosis_link = get_required_snr_se_kurtosis_on_link(mod_format_link, 4, params)
            se_link = get_required_snr_se_kurtosis_on_link(mod_format_link, 1, params)
        else:
            kurtosis_link = jnp.zeros(params.link_resources)
            se_link = jnp.ones(params.link_resources)
        bw_link = state.channel_centre_bw_array[link_index, :]
        ch_power_link = state.channel_power_array[link_index, :]
        required_slots_link = get_required_slots_on_link(bw_link, se_link, params)
        ch_centres_link = get_centre_freq_on_link(jnp.arange(params.link_resources), required_slots_link, params)

        # Calculate SNR
        P = dict(
            num_channels=params.link_resources,
            num_spans=num_spans,
            max_spans=params.max_spans,
            ref_lambda=params.ref_lambda,
            length=link_lengths,
            attenuation_i=jnp.array(params.attenuation),
            attenuation_bar_i=jnp.array(params.attenuation_bar),
            nonlinear_coeff=jnp.array(params.nonlinear_coeff),
            raman_gain_slope_i=jnp.array(params.raman_gain_slope),
            dispersion_coeff=jnp.array(params.dispersion_coeff),
            dispersion_slope=jnp.array(params.dispersion_slope),
            coherent=params.coherent,
            num_roadms=params.num_roadms,
            roadm_loss=params.roadm_loss,
            amplifier_noise_figure=params.amplifier_noise_figure.val,
            transceiver_snr=params.transceiver_snr.val,
            mod_format_correction=params.mod_format_correction,
            ch_power_w_i=ch_power_link,
            ch_centre_i=ch_centres_link*1e9,
            ch_bandwidth_i=bw_link*1e9,
            excess_kurtosis_i=kurtosis_link,
            uniform_spans=params.uniform_spans,
        )
        snr = isrs_gn_model.get_snr(**P)[0]

        return snr

    link_snr_array = jax.vmap(get_link_snr, in_axes=(0, None, None))(jnp.arange(params.num_links), state, params)
    link_snr_array = jnp.nan_to_num(link_snr_array, nan=1e-5)
    return link_snr_array

get_spectral_features(laplacian, num_features)

Compute spectral node features from symmetric normalized graph Laplacian.

Parameters:

Name	Type	Description	Default
adj		Adjacency matrix of the graph	required
num_features	int	Number of eigenvector features to extract	required

Returns:

Type	Description
ndarray	Array of shape (n_nodes, num_features) containing eigenvectors corresponding
ndarray	to the smallest non-zero eigenvalues of the graph Laplacian.

Notes

• Skips trivial eigenvectors (those with near-zero eigenvalues)
• Eigenvectors are ordered by ascending eigenvalue magnitude
• Runtime is O(n^3) - use only for small/medium graphs
• Eigenvector signs are arbitrary (may vary between runs)

Source code in xlron/environments/env_funcs.py

def get_spectral_features(laplacian: jnp.array, num_features: int) -> jnp.ndarray:
    """Compute spectral node features from symmetric normalized graph Laplacian.

    Args:
        adj: Adjacency matrix of the graph
        num_features: Number of eigenvector features to extract

    Returns:
        Array of shape (n_nodes, num_features) containing eigenvectors corresponding
        to the smallest non-zero eigenvalues of the graph Laplacian.

    Notes:
        - Skips trivial eigenvectors (those with near-zero eigenvalues)
        - Eigenvectors are ordered by ascending eigenvalue magnitude
        - Runtime is O(n^3) - use only for small/medium graphs
        - Eigenvector signs are arbitrary (may vary between runs)
    """
    n_nodes = laplacian.shape[0]
    eigenvalues, eigenvectors = jnp.linalg.eigh(laplacian)
    return eigenvectors[:, :num_features].astype(LARGE_FLOAT_DTYPE)

implement_action_rmsa_gn_model(state, action, params)

Implement action for RSA GN model. Update following arrays: - link_slot_array - link_slot_departure_array - link_snr_array - modulation_format_index_array - channel_power_array - active_path_array Args: state (EnvState): Environment state action (chex.Array): Action tuple (first is path action, second is launch_power) params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rmsa_gn_model(
        state: RSAGNModelEnvState, action: chex.Array, params: RSAGNModelEnvParams
) -> EnvState:
    """Implement action for RSA GN model. Update following arrays:
    - link_slot_array
    - link_slot_departure_array
    - link_snr_array
    - modulation_format_index_array
    - channel_power_array
    - active_path_array
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action tuple (first is path action, second is launch_power)
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_action, power_action = action
    path_action = path_action.astype(MED_INT_DTYPE)
    k_path_index, initial_slot_index = process_path_action(state, params, path_action)
    lightpath_index = get_lightpath_index(params, nodes_sd, k_path_index)
    path = get_paths(params, nodes_sd)[k_path_index]
    launch_power = get_launch_power(state, path_action, power_action, params)
    # TODO(GN MODEL) - get mod. format based on maximum reach
    mod_format_index = jax.lax.dynamic_slice(
        state.mod_format_mask, (path_action,), (1,)
    ).astype(MED_INT_DTYPE)[0]
    se = params.modulations_array.val[mod_format_index][1]
    num_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)
    # Update link_slot_array and link_slot_departure_array, then other arrays
    state = implement_path_action(state, path, initial_slot_index, num_slots)
    state = state.replace(
        path_index_array=vmap_set_path_links(state.path_index_array, path, initial_slot_index, num_slots-params.guardband, lightpath_index),
        channel_power_array=vmap_set_path_links(state.channel_power_array, path, initial_slot_index, num_slots-params.guardband, launch_power),
        modulation_format_index_array=vmap_set_path_links(state.modulation_format_index_array, path, initial_slot_index, num_slots-params.guardband, mod_format_index),
        channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, initial_slot_index, num_slots-params.guardband, params.slot_size),
    )
    # Update link_snr_array
    state = state.replace(link_snr_array=get_snr_link_array(state, params))
    # jax.debug.print("launch_power {}", launch_power, ordered=True)
    # jax.debug.print("mod_format_index {}", mod_format_index, ordered=True)
    # jax.debug.print("initial_slot_index {}", initial_slot_index, ordered=True)
    # jax.debug.print("state.mod_format_mask {}", state.mod_format_mask, ordered=True)
    # jax.debug.print("path_snr {}", get_snr_for_path(path, state.link_snr_array, params), ordered=True)
    # jax.debug.print("required_snr {}", params.modulations_array.val[mod_format_index][2] + params.snr_margin, ordered=True)
    return state

implement_action_rsa(state, action, params)

Implement action to assign slots on links.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required
action	Array	action to implement	required
params	EnvParams	environment parameters	required

Returns:

Name	Type	Description
state	EnvState	updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rsa(
        state: EnvState,
        action: chex.Array,
        params: EnvParams,
) -> EnvState:
    """Implement action to assign slots on links.

    Args:
        state: current state
        action: action to implement
        params: environment parameters

    Returns:
        state: updated state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_index, initial_slot_index = process_path_action(state, params, action)
    path = get_paths(params, nodes_sd)[path_index]
    if params.__class__.__name__ == "RWALightpathReuseEnvParams":
        state = state.replace(
            link_capacity_array=vmap_update_path_links(
                state.link_capacity_array, path, initial_slot_index, 1, requested_datarate
            )
        )
        # TODO (Dynamic-RWALR) - to support diverse requested_datarates for RWA-LR, need to update masking
        # TODO (Dynamic-RWALR) - In order to enable dynamic RWA with lightpath reuse (as opposed to just incremental loading),
        #  need to keep track of active requests OR just randomly remove connections
        #  (could do this by using the link_slot_departure array in a novel way... i.e. don't fill it with departure time but current bw)
        capacity_mask = jnp.where(state.link_capacity_array <= 0., -1., 0.)
        over_capacity_mask = jnp.where(state.link_capacity_array < 0., -1., 0.)
        total_mask = capacity_mask + over_capacity_mask
        state = state.replace(
            link_slot_array=total_mask,
            link_slot_departure_array=vmap_update_path_links(state.link_slot_departure_array, path,
                                                                       initial_slot_index, 1,
                                                                       state.current_time + state.holding_time)
        )
    else:
        se = get_paths_se(params, nodes_sd)[path_index] if params.consider_modulation_format else one
        num_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)
        state = implement_path_action(state, path, initial_slot_index, num_slots)
    return state

implement_action_rsa_gn_model(state, action, params)

Implement action for RSA GN model. Update following arrays: - link_slot_array - link_slot_departure_array - link_snr_array - modulation_format_index_array - channel_power_array - active_path_array Args: state (EnvState): Environment state action (chex.Array): Action tuple (first is path action, second is launch_power) params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rsa_gn_model(
        state: RSAGNModelEnvState, action: chex.Array, params: RSAGNModelEnvParams
) -> EnvState:
    """Implement action for RSA GN model. Update following arrays:
    - link_slot_array
    - link_slot_departure_array
    - link_snr_array
    - modulation_format_index_array
    - channel_power_array
    - active_path_array
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action tuple (first is path action, second is launch_power)
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_action, power_action = action
    path_action = path_action.astype(MED_INT_DTYPE)
    k_path_index, initial_slot_index = process_path_action(state, params, path_action)
    lightpath_index = get_lightpath_index(params, nodes_sd, k_path_index)
    path = get_paths(params, nodes_sd)[k_path_index]
    launch_power = get_launch_power(state, path_action, power_action, params)
    num_slots = required_slots(requested_datarate, 1, params.slot_size, guardband=params.guardband)
    # Update link_slot_array and link_slot_departure_array, then other arrays
    state = implement_path_action(state, path, initial_slot_index, num_slots)
    state = state.replace(
        path_index_array=vmap_set_path_links(state.path_index_array, path, initial_slot_index, num_slots-params.guardband, lightpath_index),
        channel_power_array=vmap_set_path_links(state.channel_power_array, path, initial_slot_index, num_slots-params.guardband, launch_power),
        # TODO - update this to use separate arrays to track channel centres and bandwidths and update with bandwidth (that may or may not equal slot size)
        channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, initial_slot_index, num_slots-params.guardband, params.slot_size),
    )
    if params.monitor_active_lightpaths:
        state = state.replace(
            active_lightpaths_array=update_active_lightpaths_array(state, lightpath_index, initial_slot_index, num_slots-params.guardband),
            active_lightpaths_array_departure=update_active_lightpaths_array_departure(state, -state.current_time-state.holding_time),
        )
        # No need to check SNR until end of episode
        return state
    # Update link_snr_array
    state = state.replace(link_snr_array=get_snr_link_array(state, params))
    return state

implement_action_rwalr(state, action, params)

For use in RWALightpathReuseEnv. Update link_slot_array and link_slot_departure_array to reflect new lightpath assignment. Update link_capacity_array with new capacity if lightpath is available. Undo link_capacity_update if over capacity.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
action	Array	Action array	required
params	EnvParams	Environment parameters	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rwalr(state: EnvState, action: chex.Array, params: EnvParams) -> EnvState:
    """For use in RWALightpathReuseEnv.
    Update link_slot_array and link_slot_departure_array to reflect new lightpath assignment.
    Update link_capacity_array with new capacity if lightpath is available.
    Undo link_capacity_update if over capacity.

    Args:
        state: Environment state
        action: Action array
        params: Environment parameters

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_index, initial_slot_index = process_path_action(state, params, action)
    path = get_paths(params, nodes_sd)[path_index]
    lightpath_available_check, lightpath_existing_check, curr_lightpath_capacity, lightpath_index = (
        check_lightpath_available_and_existing(state, params, action)
    )
    # Get path capacity - request
    lightpath_capacity = jax.lax.cond(
        lightpath_existing_check,
        lambda x: curr_lightpath_capacity - requested_datarate,  # Subtract requested_datarate from current lightpath
        lambda x: jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.path_capacity_array, x, 1)) - requested_datarate,  # Get initial capacity of lightpath - request
        lightpath_index
    )
    # Update link_capacity_array with new capacity if lightpath is available
    state = jax.lax.cond(
        lightpath_available_check,
        lambda x: x.replace(
            link_capacity_array=vmap_set_path_links(
                state.link_capacity_array, path, initial_slot_index, 1, lightpath_capacity
            ),
            path_index_array=vmap_set_path_links(
                state.path_index_array, path, initial_slot_index, 1, lightpath_index
            ),
        ),
        lambda x: x,
        state
    )
    capacity_mask = jnp.where(state.link_capacity_array <= 0., -1., 0.)
    over_capacity_mask = jnp.where(state.link_capacity_array < 0., -1., 0.)
    # Undo link_capacity_update if over capacity
    # N.B. this will fail if requested capacity is greater than total original capacity of lightpath
    lightpath_capacity_before_action = jax.lax.cond(
        lightpath_existing_check,
        lambda x: curr_lightpath_capacity,  # Subtract requested_datarate from current lightpath
        lambda x: 1e6,  # Empty slots have high capacity (1e6)
        # Get initial capacity of lightpath - request
        None,
    )
    state = state.replace(
        link_capacity_array=jnp.where(over_capacity_mask == -1, lightpath_capacity_before_action, state.link_capacity_array)
    )
    # Total mask will be 0 if space still available, -1 if capacity is zero or -2 if over capacity
    total_mask = capacity_mask + over_capacity_mask
    # Update link_slot_array and link_slot_departure_array
    state = state.replace(
        link_slot_array=total_mask,
        link_slot_departure_array=vmap_update_path_links(state.link_slot_departure_array, path,
                                                                 initial_slot_index, 1,
                                                                 state.current_time + state.holding_time)
    )
    return state

implement_node_action(state, s_node, d_node, s_request, d_request, n=2)

Update node capacity, node resource and node departure arrays

Parameters:

Name	Type	Description	Default
state	State	current state	required
s_node	int	source node	required
d_node	int	destination node	required
s_request	int	source node request	required
d_request	int	destination node request	required
n	int	number of nodes to implement. Defaults to 2.	2

Returns:

Name	Type	Description
State	EnvState	updated state

Source code in xlron/environments/env_funcs.py

def implement_node_action(state: EnvState, s_node: chex.Array, d_node: chex.Array, s_request: chex.Array, d_request: chex.Array, n=2) -> EnvState:
    """Update node capacity, node resource and node departure arrays

    Args:
        state (State): current state
        s_node (int): source node
        d_node (int): destination node
        s_request (int): source node request
        d_request (int): destination node request
        n (int, optional): number of nodes to implement. Defaults to 2.

    Returns:
        State: updated state
    """
    node_indices = jnp.arange(state.node_capacity_array.shape[0])

    curr_selected_nodes = jnp.zeros(state.node_capacity_array.shape[0])
    # d_request -ve so that selected node is +ve (so that argmin works correctly for node resource array update)
    # curr_selected_nodes is N x 1 array, with requested node resources at index of selected node
    curr_selected_nodes = update_node_array(node_indices, curr_selected_nodes, d_node, -d_request)
    curr_selected_nodes = jax.lax.cond(n == 2, lambda x: update_node_array(*x), lambda x: x[1], (node_indices, curr_selected_nodes, s_node, -s_request))

    node_capacity_array = state.node_capacity_array - curr_selected_nodes

    node_resource_array = vmap_update_node_resources(state.node_resource_array, curr_selected_nodes)

    node_departure_array = vmap_update_node_departure(state.node_departure_array, curr_selected_nodes, -state.current_time-state.holding_time)

    state = state.replace(
        node_capacity_array=node_capacity_array,
        node_resource_array=node_resource_array,
        node_departure_array=node_departure_array
    )

    return state

implement_path_action(state, path, initial_slot_index, num_slots)

Update link-slot and link-slot departure arrays. Times are set to negative until turned positive by finalisation (after checks).

Parameters:

Name	Type	Description	Default
state	State	current state	required
path	int	path to implement	required
initial_slot_index	int	initial slot index	required
num_slots	int	number of slots to implement	required

Source code in xlron/environments/env_funcs.py

def implement_path_action(state: EnvState, path: chex.Array, initial_slot_index: chex.Array, num_slots: chex.Array) -> EnvState:
    """Update link-slot and link-slot departure arrays.
    Times are set to negative until turned positive by finalisation (after checks).

    Args:
        state (State): current state
        path (int): path to implement
        initial_slot_index (int): initial slot index
        num_slots (int): number of slots to implement
    """
    state = state.replace(
        link_slot_array=vmap_update_path_links(state.link_slot_array, path, initial_slot_index, num_slots, one),
        link_slot_departure_array=vmap_update_path_links(state.link_slot_departure_array, path, initial_slot_index, num_slots, state.current_time+state.holding_time)
    )
    return state

implement_vone_action(state, action, total_actions, remaining_actions, params)

Implement action to assign nodes (1, 2, or 0 nodes assigned per action) and assign slots and links for lightpath.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required
action	Array	action to implement (node, node, path_slot_action)	required
total_actions	Scalar	total number of actions to implement for current request	required
remaining_actions	Scalar	remaining actions to implement	required
k		number of paths to consider	required
N		number of nodes to assign	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(4,))
def implement_vone_action(
        state: EnvState,
        action: chex.Array,
        total_actions: chex.Scalar,
        remaining_actions: chex.Scalar,
        params: EnvParams,
):
    """Implement action to assign nodes (1, 2, or 0 nodes assigned per action) and assign slots and links for lightpath.

    Args:
        state: current state
        action: action to implement (node, node, path_slot_action)
        total_actions: total number of actions to implement for current request
        remaining_actions: remaining actions to implement
        k: number of paths to consider
        N: number of nodes to assign

    Returns:
        state: updated state
    """
    request = jax.lax.dynamic_slice(state.request_array[0], ((remaining_actions-1)*2, ), (3, ))
    node_request_s = jax.lax.dynamic_slice(request, (2, ), (1, ))
    requested_datarate = jax.lax.dynamic_slice(request, (1,), (1,))
    node_request_d = jax.lax.dynamic_slice(request, (0, ), (1, ))
    nodes = action[::2]
    path_index, initial_slot_index = process_path_action(state, params, action[1])
    path = get_paths(params, nodes)[path_index]
    se = get_paths_se(params, nodes)[path_index] if params.consider_modulation_format else jnp.array([1])
    num_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)

    # jax.debug.print("state.request_array {}", state.request_array, ordered=True)
    # jax.debug.print("path {}", path, ordered=True)
    # jax.debug.print("slots {}", jnp.max(jnp.where(path.reshape(-1,1) == 1, state.link_slot_array, jnp.zeros(params.num_links).reshape(-1,1)), axis=0), ordered=True)
    # jax.debug.print("path_index {}", path_index, ordered=True)
    # jax.debug.print("initial_slot_index {}", initial_slot_index, ordered=True)
    # jax.debug.print("requested_datarate {}", requested_datarate, ordered=True)
    # jax.debug.print("request {}", request, ordered=True)
    # jax.debug.print("se {}", se, ordered=True)
    # jax.debug.print("num_slots {}", num_slots, ordered=True)

    n_nodes = jax.lax.cond(
        total_actions == remaining_actions,
        lambda x: 2, lambda x: 1,
        (total_actions, remaining_actions))
    path_action_only_check = path_action_only(state.request_array[1], state.action_counter, remaining_actions)

    state = jax.lax.cond(
        path_action_only_check,
        lambda x: x[0],
        lambda x: implement_node_action(x[0], x[1], x[2], x[3], x[4], n=x[5]),
        (state, nodes[0], nodes[1], node_request_s, node_request_d, n_nodes)
    )

    state = implement_path_action(state, path, initial_slot_index, num_slots)

    return state

init_action_counter()

Initialize action counter. First index is num unique nodes, second index is total steps, final is remaining steps until completion of request. Only used in VONE environments.

Source code in xlron/environments/env_funcs.py

def init_action_counter():
    """Initialize action counter.
    First index is num unique nodes, second index is total steps, final is remaining steps until completion of request.
    Only used in VONE environments.
    """
    return jnp.zeros(3, dtype=MED_INT_DTYPE)

init_action_history(params)

Initialize action history

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_action_history(params: EnvParams):
    """Initialize action history"""
    return jnp.full(params.max_edges*2+1, -1, dtype=LARGE_FLOAT_DTYPE)

init_active_lightpaths_array(params)

Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array. M is MIN(max_requests, num_links * link_resources / min_slots). min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

Parameters:

Name	Type	Description	Default
params	RSAGNModelEnvParams	Environment parameters	required

Returns: jnp.array: Active path array (default value -1, empty path)

Source code in xlron/environments/env_funcs.py

def init_active_lightpaths_array(params: RSAGNModelEnvParams):
    """Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array.
    M is MIN(max_requests, num_links * link_resources / min_slots).
    min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

    Args:
        params (RSAGNModelEnvParams): Environment parameters
    Returns:
        jnp.array: Active path array (default value -1, empty path)
    """
    total_slots = params.num_links * params.link_resources  # total slots on networks
    min_slots = jnp.max(params.values_bw.val) / params.slot_size  # minimum number of slots required for lightpath
    return jnp.full((int(total_slots / min_slots), 3), -1, dtype=LARGE_INT_DTYPE)

init_active_lightpaths_array_departure(params)

Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array. M is MIN(max_requests, num_links * link_resources / min_slots). min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

Parameters:

Name	Type	Description	Default
params	RSAGNModelEnvParams	Environment parameters	required

Returns: jnp.array: Active path array (default value -1, empty path)

Source code in xlron/environments/env_funcs.py

def init_active_lightpaths_array_departure(params: RSAGNModelEnvParams):
    """Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array.
    M is MIN(max_requests, num_links * link_resources / min_slots).
    min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

    Args:
        params (RSAGNModelEnvParams): Environment parameters
    Returns:
        jnp.array: Active path array (default value -1, empty path)
    """
    total_slots = params.num_links * params.link_resources  # total slots on networks
    min_slots = jnp.max(params.values_bw.val) / params.slot_size  # minimum number of slots required for lightpath
    return jnp.full((int(total_slots / min_slots), 3), 0., dtype=SMALL_FLOAT_DTYPE)

init_active_path_array(params)

Initialise active path array. Stores details of full path utilised by lightpath on each frequency slot. Args: params (EnvParams): Environment parameters Returns: jnp.array: Active path array

Source code in xlron/environments/env_funcs.py

def init_active_path_array(params: EnvParams):
    """Initialise active path array. Stores details of full path utilised by lightpath on each frequency slot.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Active path array
    """
    return jnp.full((params.num_links, params.link_resources, params.num_links), -1, dtype=MED_INT_DTYPE)

init_channel_centre_bw_array(params)

Initialise channel centre array. Args: params (EnvParams): Environment parameters Returns: jnp.array: Channel centre array

Source code in xlron/environments/env_funcs.py

def init_channel_centre_bw_array(params: EnvParams):
    """Initialise channel centre array.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Channel centre array
    """
    return jnp.full((params.num_links, params.link_resources), 0., dtype=LARGE_FLOAT_DTYPE)

init_channel_power_array(params)

Initialise channel power array.

Parameters:

Name	Type	Description	Default
params	EnvParams	Environment parameters	required

Returns: jnp.array: Channel power array

Source code in xlron/environments/env_funcs.py

def init_channel_power_array(params: EnvParams):
    """Initialise channel power array.

    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Channel power array
    """
    return jnp.full((params.num_links, params.link_resources), 0., dtype=LARGE_FLOAT_DTYPE)

init_graph_tuple(state, params, adj, exclude_source_dest=False)

Initialise graph tuple for use with Jraph GNNs. Args: state (EnvState): Environment state params (EnvParams): Environment parameters adj (jnp.array): Adjacency matrix of the graph Returns: jraph.GraphsTuple: Graph tuple

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 3))
def init_graph_tuple(state: EnvState, params: EnvParams, adj: jnp.array, exclude_source_dest: bool=False) -> jraph.GraphsTuple:
    """Initialise graph tuple for use with Jraph GNNs.
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        adj (jnp.array): Adjacency matrix of the graph
    Returns:
        jraph.GraphsTuple: Graph tuple
    """
    senders = params.edges.val.T[0]
    receivers = params.edges.val.T[1]

    # Get source and dest from request array
    source_dest, datarate = read_rsa_request(state.request_array)
    # Global feature is normalised data rate of current request
    globals = jnp.array([datarate / jnp.max(params.values_bw.val)], dtype=LARGE_FLOAT_DTYPE)

    if exclude_source_dest:
        source_dest_features = jnp.zeros((params.num_nodes, 2), dtype=LARGE_FLOAT_DTYPE)
    else:
        source, dest = source_dest[0], source_dest[2]
        # One-hot encode source and destination (2 additional features)
        source_dest_features = jnp.zeros((params.num_nodes, 2), dtype=LARGE_FLOAT_DTYPE)
        source_dest_features = source_dest_features.at[source.astype(MED_INT_DTYPE), 0].set(1)
        source_dest_features = source_dest_features.at[dest.astype(MED_INT_DTYPE), 1].set(-1)

    spectral_features = get_spectral_features(adj, num_features=3)

    # For dynamic traffic, edge_features are normalised remaining holding time instead of link_slot_array
    holding_time_edge_features = state.link_slot_departure_array / params.mean_service_holding_time

    if params.__class__.__name__ in ["RSAGNModelEnvParams", "RMSAGNModelEnvParams"]:
        # Normalize by max parameters (converted to linear units)
        max_power = isrs_gn_model.from_dbm(params.max_power)
        normalized_power = jnp.round(state.channel_power_array / max_power, 3)
        max_snr = isrs_gn_model.from_db(params.max_snr)
        normalized_snr = jnp.round(state.link_snr_array / max_snr, 3)
        edge_features = jnp.stack([normalized_snr, normalized_power], axis=-1)
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)
    elif params.__class__.__name__ == "VONEEnvParams":
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        node_features = getattr(state, "node_capacity_array", jnp.zeros(params.num_nodes, dtype=LARGE_FLOAT_DTYPE))
        node_features = node_features.reshape(-1, 1)
        node_features = jnp.concatenate([node_features, spectral_features, source_dest_features], axis=-1)
    else:
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        # [n_edges] or [n_edges, ...]
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)

    if params.disable_node_features:
        node_features = jnp.zeros((1,), dtype=LARGE_FLOAT_DTYPE)

    # Handle undirected graphs (duplicate edges after normalization)
    if not params.directed_graph:
        senders_ = jnp.concatenate([senders, receivers])
        receivers = jnp.concatenate([receivers, senders])
        senders = senders_
        edge_features = jnp.repeat(edge_features, 2, axis=0)

    return jraph.GraphsTuple(
        nodes=node_features,
        edges=edge_features,
        senders=senders,
        receivers=receivers,
        n_node=jnp.reshape(params.num_nodes, (1,)),
        n_edge=jnp.reshape(len(senders), (1,)),
        globals=globals,
    )

init_link_capacity_array(params)

Initialise link capacity array. Represents available data rate for lightpath on each link. Default is high value (1e6) for unoccupied slots. Once lightpath established, capacity is determined by corresponding entry in path capacity array.

Source code in xlron/environments/env_funcs.py

def init_link_capacity_array(params):
    """Initialise link capacity array. Represents available data rate for lightpath on each link.
    Default is high value (1e6) for unoccupied slots. Once lightpath established, capacity is determined by
    corresponding entry in path capacity array."""
    return jnp.full((params.num_links, params.link_resources), 1e6)

init_link_length_array(graph)

Initialise link length array. Args: graph (nx.Graph): NetworkX graph Returns:

Source code in xlron/environments/env_funcs.py

def init_link_length_array(graph: nx.Graph) -> chex.Array:
    """Initialise link length array.
    Args:
        graph (nx.Graph): NetworkX graph
    Returns:

    """
    link_lengths = []
    for edge in sorted(graph.edges):
        link_lengths.append(graph.edges[edge]["weight"])
    return jnp.array(link_lengths, dtype=MED_INT_DTYPE)

init_link_length_array_gn_model(graph, max_span_length, max_spans)

Initialise link length array for environements that use GN model of physical layer. We assume each link has spans of equal length.

Parameters:

Name	Type	Description	Default
graph	Graph	NetworkX graph	required

Returns: jnp.array: Link length array (L x max_spans)

Source code in xlron/environments/env_funcs.py

def init_link_length_array_gn_model(graph: nx.Graph, max_span_length: int,  max_spans: int) -> chex.Array:
    """Initialise link length array for environements that use GN model of physical layer.
    We assume each link has spans of equal length.

    Args:
        graph (nx.Graph): NetworkX graph
    Returns:
        jnp.array: Link length array (L x max_spans)
    """
    link_lengths = []
    directed = graph.is_directed()
    graph = graph.to_undirected()
    edges = sorted(graph.edges)
    for edge in edges:
        link_lengths.append(graph.edges[edge]["weight"])
    if directed:
        for edge in edges:
            link_lengths.append(graph.edges[edge]["weight"])
    span_length_array = []
    for length in link_lengths:
        num_spans = math.ceil(length / max_span_length)
        avg_span_length = length / num_spans
        span_lengths = [avg_span_length] * num_spans
        span_lengths.extend([0] * (max_spans - num_spans))
        span_length_array.append(span_lengths)
    return jnp.array(span_length_array, dtype=MED_INT_DTYPE)

init_link_slot_array(params)

Initialize empty (all zeroes) link-slot array. 0 means slot is free, -1 means occupied. Args: params (EnvParams): Environment parameters Returns: jnp.array: Link slot array (E x S) where E is number of edges and S is number of slots

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_link_slot_array(params: EnvParams):
    """Initialize empty (all zeroes) link-slot array. 0 means slot is free, -1 means occupied.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Link slot array (E x S) where E is number of edges and S is number of slots"""
    return jnp.zeros((params.num_links, params.link_resources), dtype=LARGE_FLOAT_DTYPE)

init_link_slot_mask(params, agg=1)

Initialize link mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0, 1))
def init_link_slot_mask(params: EnvParams, agg: int = 1):
    """Initialize link mask"""
    return jnp.ones(params.k_paths*math.ceil(params.link_resources / agg), dtype=LARGE_FLOAT_DTYPE)

init_link_snr_array(params)

Initialise signal-to-noise ratio (SNR) array. Args: params (EnvParams): Environment parameters Returns: jnp.array: SNR array

Source code in xlron/environments/env_funcs.py

def init_link_snr_array(params: EnvParams):
    """Initialise signal-to-noise ratio (SNR) array.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: SNR array
    """
    # The SNR is kept in linear units to allow summation of 1/SNR across links
    return jnp.full((params.num_links, params.link_resources), -1e5, dtype=LARGE_FLOAT_DTYPE)

init_mod_format_mask(params)

Initialize link mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_mod_format_mask(params: EnvParams):
    """Initialize link mask"""
    return jnp.full((params.k_paths*params.link_resources,), -1.0, dtype=LARGE_FLOAT_DTYPE)

init_modulation_format_index_array(params)

Initialise modulation format index array. Args: params (EnvParams): Environment parameters Returns: jnp.array: Modulation format index array

Source code in xlron/environments/env_funcs.py

def init_modulation_format_index_array(params: EnvParams):
    """Initialise modulation format index array.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Modulation format index array
    """
    return jnp.full((params.num_links, params.link_resources), -1, dtype=MED_INT_DTYPE)  # -1 so that highest order is assumed (closest to Gaussian)

init_modulations_array(modulations_filepath=None)

Initialise array of maximum spectral efficiency for modulation format on path.

Parameters:

Name	Type	Description	Default
modulations_filepath	str	Path to CSV file containing modulation formats. Defaults to None.	None

Returns: jnp.array: Array of maximum spectral efficiency for modulation format on path. First two columns are maximum path length and spectral efficiency.

Source code in xlron/environments/env_funcs.py

def init_modulations_array(modulations_filepath: str = None):
    """Initialise array of maximum spectral efficiency for modulation format on path.

    Args:
        modulations_filepath (str, optional): Path to CSV file containing modulation formats. Defaults to None.
    Returns:
        jnp.array: Array of maximum spectral efficiency for modulation format on path.
        First two columns are maximum path length and spectral efficiency.
    """
    f = pathlib.Path(modulations_filepath) if modulations_filepath else (
            pathlib.Path(__file__).parents[1].absolute() / "data" / "modulations" / "modulations.csv")
    modulations = np.genfromtxt(f, delimiter=',')
    # Drop empty first row (headers) and column (name)
    modulations = modulations[1:, 1:]
    return jnp.array(modulations, dtype=LARGE_FLOAT_DTYPE)

init_node_capacity_array(params)

Initialize node array with uniform resources. Args: params (EnvParams): Environment parameters Returns: jnp.array: Node capacity array (N x 1) where N is number of nodes

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_node_capacity_array(params: EnvParams):
    """Initialize node array with uniform resources.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Node capacity array (N x 1) where N is number of nodes"""
    return jnp.array([params.node_resources] * params.num_nodes, dtype=MED_INT_DTYPE)

init_node_mask(params)

Initialize node mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_node_mask(params: EnvParams):
    """Initialize node mask"""
    return jnp.ones(params.num_nodes, dtype=LARGE_FLOAT_DTYPE)

init_node_resource_array(params)

Array to track node resources occupied by virtual nodes

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_node_resource_array(params: EnvParams):
    """Array to track node resources occupied by virtual nodes"""
    return jnp.zeros((params.num_nodes, params.node_resources), dtype=LARGE_FLOAT_DTYPE)

init_path_capacity_array(link_length_array, path_link_array, min_request=1, scale_factor=1.0, alpha=0.0002, NF=4.5, B=10000000000000.0, R_s=100000000000.0, beta_2=-2.17e-26, gamma=0.0012, L_s=100000.0, lambda0=1.55e-06)

Calculated from Nevin paper: https://api.repository.cam.ac.uk/server/api/core/bitstreams/b80e7a9c-a86b-4b30-a6d6-05017c60b0c8/content

Parameters:

Name	Type	Description	Default
link_length_array	Array	Array of link lengths	required
path_link_array	Array	Array of links on paths	required
min_request	int	Minimum data rate request size. Defaults to 100 GBps.	1
scale_factor	float	Scale factor for link capacity. Defaults to 1.0.	1.0
alpha	float	Fibre attenuation coefficient. Defaults to 0.2e-3 /m	0.0002
NF	float	Amplifier noise figure. Defaults to 4.5 dB.	4.5
B	float	Total modulated bandwidth. Defaults to 10e12 Hz.	10000000000000.0
R_s	float	Symbol rate. Defaults to 100e9 Baud.	100000000000.0
beta_2	float	Dispersion parameter. Defaults to -21.7e-27 s^2/m.	-2.17e-26
gamma	float	Nonlinear coefficient. Defaults to 1.2e-3 /W/m.	0.0012
L_s	float	Span length. Defaults to 100e3 m.	100000.0
lambda0	float	Wavelength. Defaults to 1550e-9 m.	1.55e-06

Returns:

Type	Description
Array	chex.Array: Array of link capacities in Gbps

Source code in xlron/environments/env_funcs.py

def init_path_capacity_array(
        link_length_array: chex.Array,
        path_link_array: chex.Array,
        min_request=1,  # Minimum data rate request size
        scale_factor=1.0,  # Scale factor for link capacity
        alpha=0.2e-3,  # Fibre attenuation coefficient
        NF=4.5,  # Amplifier noise figure
        B=10e12,  # Total modulated bandwidth
        R_s=100e9,  # Symbol rate
        beta_2=-21.7e-27,  # Dispersion parameter
        gamma=1.2e-3,  # Nonlinear coefficient
        L_s=100e3,  # span length
        lambda0=1550e-9,  # Wavelength
) -> chex.Array:
    """Calculated from Nevin paper:
    https://api.repository.cam.ac.uk/server/api/core/bitstreams/b80e7a9c-a86b-4b30-a6d6-05017c60b0c8/content

    Args:
        link_length_array (chex.Array): Array of link lengths
        path_link_array (chex.Array): Array of links on paths
        min_request (int, optional): Minimum data rate request size. Defaults to 100 GBps.
        scale_factor (float, optional): Scale factor for link capacity. Defaults to 1.0.
        alpha (float, optional): Fibre attenuation coefficient. Defaults to 0.2e-3 /m
        NF (float, optional): Amplifier noise figure. Defaults to 4.5 dB.
        B (float, optional): Total modulated bandwidth. Defaults to 10e12 Hz.
        R_s (float, optional): Symbol rate. Defaults to 100e9 Baud.
        beta_2 (float, optional): Dispersion parameter. Defaults to -21.7e-27 s^2/m.
        gamma (float, optional): Nonlinear coefficient. Defaults to 1.2e-3 /W/m.
        L_s (float, optional): Span length. Defaults to 100e3 m.
        lambda0 (float, optional): Wavelength. Defaults to 1550e-9 m.

    Returns:
        chex.Array: Array of link capacities in Gbps
    """
    path_length_array = jnp.dot(path_link_array, link_length_array)
    path_capacity_array = calculate_path_capacity(
        path_length_array,
        min_request=min_request,
        scale_factor=scale_factor,
        alpha=alpha,
        NF=NF,
        B=B,
        R_s=R_s,
        beta_2=beta_2,
        gamma=gamma,
        L_s=L_s,
        lambda0=lambda0,
    )
    return path_capacity_array.astype(MED_INT_DTYPE)

init_path_index_array(params)

Initialise path index array. Represents index of lightpath occupying each slot.

Source code in xlron/environments/env_funcs.py

def init_path_index_array(params):
    """Initialise path index array. Represents index of lightpath occupying each slot."""
    return jnp.full((params.num_links, params.link_resources), -1)

init_path_length_array(path_link_array, graph)

Initialise path length array.

Parameters:

Name	Type	Description	Default
path_link_array	Array	Path-link array	required
graph	Graph	NetworkX graph	required

Returns: chex.Array: Path length array

Source code in xlron/environments/env_funcs.py

def init_path_length_array(path_link_array: chex.Array, graph: nx.Graph) -> chex.Array:
    """Initialise path length array.

    Args:
        path_link_array (chex.Array): Path-link array
        graph (nx.Graph): NetworkX graph
    Returns:
        chex.Array: Path length array
    """
    link_length_array = init_link_length_array(graph)
    path_lengths = jnp.dot(path_link_array, link_length_array)
    return path_lengths

init_path_link_array(graph, k, disjoint=False, weight='weight', directed=False, modulations_array=None, rwa_lr=False, scale_factor=1.0, path_snr=False)

Initialise path-link array. Each path is defined by a link utilisation array (one row in the path-link array). 1 indicates link corresponding to index is used, 0 indicates not used.

Parameters:

Name	Type	Description	Default
graph	Graph	NetworkX graph	required
k	int	Number of paths	required
disjoint	bool	Whether to use edge-disjoint paths. Defaults to False.	False
weight	str	Sort paths by edge attribute. Defaults to "weight".	'weight'
directed	bool	Whether graph is directed. Defaults to False.	False
modulations_array	Array	Array of maximum spectral efficiency for modulation format on path. Defaults to None.	None
rwa_lr	bool	Whether the environment is RWA with lightpath reuse (affects path ordering).	False
path_snr	bool	If GN model is used, include extra row of zeroes for unutilised paths	False

Returns:

Type	Description
Array	chex.Array: Path-link array (N(N-1)*k x E) where N is number of nodes, E is number of edges, k is number of shortest paths

Source code in xlron/environments/env_funcs.py

def init_path_link_array(
        graph: nx.Graph,
        k: int,
        disjoint: bool = False,
        weight: str = "weight",
        directed: bool = False,
        modulations_array: chex.Array = None,
        rwa_lr: bool = False,
        scale_factor: float = 1.0,
        path_snr: bool = False,
) -> chex.Array:
    """Initialise path-link array.
    Each path is defined by a link utilisation array (one row in the path-link array).
    1 indicates link corresponding to index is used, 0 indicates not used.

    Args:
        graph (nx.Graph): NetworkX graph
        k (int): Number of paths
        disjoint (bool, optional): Whether to use edge-disjoint paths. Defaults to False.
        weight (str, optional): Sort paths by edge attribute. Defaults to "weight".
        directed (bool, optional): Whether graph is directed. Defaults to False.
        modulations_array (chex.Array, optional): Array of maximum spectral efficiency for modulation format on path. Defaults to None.
        rwa_lr (bool, optional): Whether the environment is RWA with lightpath reuse (affects path ordering).
        path_snr (bool, optional): If GN model is used, include extra row of zeroes for unutilised paths
        to ensure correct SNR calculation for empty paths (path index -1).

    Returns:
        chex.Array: Path-link array (N(N-1)*k x E) where N is number of nodes, E is number of edges, k is number of shortest paths
    """
    def get_k_shortest_paths(g, source, target, k, weight):
        return list(
            islice(nx.shortest_simple_paths(g, source, target, weight=weight), k)
        )

    def get_k_disjoint_shortest_paths(g, source, target, k, weight):
        k_paths_disjoint_unsorted = list(nx.edge_disjoint_paths(g, source, target))
        k_paths_shortest = get_k_shortest_paths(g, source, target, k, weight=weight)

        # Keep disjoint paths and add unique shortest paths until k paths reached
        disjoint_ids = [tuple(path) for path in k_paths_disjoint_unsorted]
        k_paths = k_paths_disjoint_unsorted
        for path in k_paths_shortest:
            if tuple(path) not in disjoint_ids:
                k_paths.append(path)
        k_paths = k_paths[:k]
        return k_paths

    paths = []
    edges = sorted(graph.edges)

    # Get the k-shortest paths for each node pair
    k_path_collections = []
    get_paths = get_k_disjoint_shortest_paths if disjoint else get_k_shortest_paths
    for node_pair in combinations(graph.nodes, 2):

        k_paths = get_paths(graph, node_pair[0], node_pair[1], k, weight=weight)
        k_path_collections.append(k_paths)

    if directed:  # Get paths in reverse direction
        for node_pair in combinations(graph.nodes, 2):
            k_paths_rev = get_paths(graph, node_pair[1], node_pair[0], k, weight=weight)
            k_path_collections.append(k_paths_rev)

    # Sort the paths for each node pair
    max_missing_paths = 0
    for k_paths in k_path_collections:

        source, dest = k_paths[0][0], k_paths[0][-1]

        # Sort the paths by # of hops then by length, or just length
        path_lengths = [nx.path_weight(graph, path, weight='weight') for path in k_paths]
        path_num_links = [len(path) - 1 for path in k_paths]

        # Get maximum spectral efficiency for modulation format on path
        if modulations_array is not None and rwa_lr is not True:
            se_of_path = []
            modulations_array = modulations_array[::-1]
            for length in path_lengths:
                for modulation in modulations_array:
                    if length <= modulation[0]:
                        se_of_path.append(modulation[1])
                        break
            # Sorting by the num_links/se instead of just path length is observed to improve performance
            path_weighting = [num_links/se for se, num_links in zip(se_of_path, path_num_links)]
        elif rwa_lr:
            path_capacity = [float(calculate_path_capacity(path_length, scale_factor=scale_factor)) for path_length in path_lengths]
            path_weighting = [num_links/path_capacity for num_links, path_capacity in zip(path_num_links, path_capacity)]
        elif weight is None:
            path_weighting = path_num_links
        else:
            path_weighting = path_lengths

        # if less then k unique paths, add empty paths
        empty_path = [0] * len(graph.edges)
        num_missing_paths = k - len(k_paths)
        max_missing_paths = max(max_missing_paths, num_missing_paths)
        k_paths = k_paths + [empty_path] * num_missing_paths
        path_weighting = path_weighting + [1e6] * num_missing_paths
        path_lengths = path_lengths + [1e6] * num_missing_paths

        # Sort by number of links then by length (or just by length if weight is specified)
        unsorted_paths = zip(k_paths, path_weighting, path_lengths)
        k_paths_sorted = [(source, dest, weighting, path) for path, weighting, _ in sorted(unsorted_paths, key=lambda x: (x[1], 1/x[2]) if weight is None else x[2])]

        # Keep only first k paths
        k_paths_sorted = k_paths_sorted[:k]

        prev_link_usage = empty_path
        for k_path in k_paths_sorted:
            k_path = k_path[-1]
            link_usage = [0]*len(graph.edges)  # Initialise empty path
            if sum(k_path) == 0:
                link_usage = prev_link_usage
            else:
                for i in range(len(k_path)-1):
                    s, d = k_path[i], k_path[i + 1]
                    for edge_index, edge in enumerate(edges):
                        condition = (edge[0] == s and edge[1] == d) if directed else \
                            ((edge[0] == s and edge[1] == d) or (edge[0] == d and edge[1] == s))
                        if condition:
                            link_usage[edge_index] = 1
            path = link_usage
            prev_link_usage = link_usage
            paths.append(path)

    # If using GN model, add extra row of zeroes for empty paths for SNR calculation
    if path_snr:
        empty_path = [0] * len(graph.edges)
        paths.append(empty_path)

    return jnp.array(paths, dtype=SMALL_INT_DTYPE)

init_path_se_array(path_length_array, modulations_array)

Initialise array of maximum spectral efficiency for highest-order modulation format on path.

Parameters:

Name	Type	Description	Default
path_length_array	array	Array of path lengths	required
modulations_array	array	Array of maximum spectral efficiency for modulation format on path	required

Returns:

Type	Description
	jnp.array: Array of maximum spectral efficiency for on path

Source code in xlron/environments/env_funcs.py

def init_path_se_array(path_length_array, modulations_array):
    """Initialise array of maximum spectral efficiency for highest-order modulation format on path.

    Args:
        path_length_array (jnp.array): Array of path lengths
        modulations_array (jnp.array): Array of maximum spectral efficiency for modulation format on path

    Returns:
        jnp.array: Array of maximum spectral efficiency for on path
    """
    se_list = []
    # Flip the modulation array so that the shortest path length is first
    modulations_array = modulations_array[::-1]
    for length in path_length_array:
        for modulation in modulations_array:
            if length <= modulation[0]:
                se_list.append(modulation[1])
                break
    return jnp.array(se_list, dtype=SMALL_INT_DTYPE)

init_rsa_request_array()

Initialize request array

Source code in xlron/environments/env_funcs.py

def init_rsa_request_array():
    """Initialize request array"""
    return jnp.zeros(3, dtype=MED_INT_DTYPE)

init_traffic_matrix(key, params)

Initialize traffic matrix. Allows for random traffic matrix or uniform traffic matrix. Source-dest traffic requests are sampled probabilistically from the resulting traffic matrix.

Parameters:

Name	Type	Description	Default
key	PRNGKey	PRNG key	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
	jnp.array: Traffic matrix

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def init_traffic_matrix(key: chex.PRNGKey, params: EnvParams):
    """Initialize traffic matrix. Allows for random traffic matrix or uniform traffic matrix.
    Source-dest traffic requests are sampled probabilistically from the resulting traffic matrix.

    Args:
        key (chex.PRNGKey): PRNG key
        params (EnvParams): Environment parameters

    Returns:
        jnp.array: Traffic matrix
    """
    if params.random_traffic:
        traffic_matrix = jax.random.uniform(key, shape=(params.num_nodes, params.num_nodes), dtype=SMALL_FLOAT_DTYPE)
    else:
        traffic_matrix = jnp.ones((params.num_nodes, params.num_nodes), dtype=SMALL_FLOAT_DTYPE)
    diag_elements = jnp.diag_indices_from(traffic_matrix)
    # Set main diagonal to zero so no requests from node to itself
    traffic_matrix = traffic_matrix.at[diag_elements].set(0)
    traffic_matrix = normalise_traffic_matrix(traffic_matrix)
    return traffic_matrix

init_transceiver_amplifier_noise_arrays(link_resources, ref_lambda, slot_size, noise_data_filepath=None)

Initialise transceiver and amplifier noise arrays. Args: link_resources (int): Number of link resources ref_lambda (float): Reference wavelength slot_size (float): Slot size noise_data_filepath (str, optional): Path to CSV file containing modulation formats. Defaults to None. Returns: Tuple[chex.Array, chex.Array]: Transceiver noise array, Amplifier noise array

Source code in xlron/environments/env_funcs.py

def init_transceiver_amplifier_noise_arrays(
        link_resources: int,
        ref_lambda: float,
        slot_size: float,
        noise_data_filepath: str = None
) -> Tuple[chex.Array, chex.Array]:
    """Initialise transceiver and amplifier noise arrays.
    Args:
        link_resources (int): Number of link resources
        ref_lambda (float): Reference wavelength
        slot_size (float): Slot size
        noise_data_filepath (str, optional): Path to CSV file containing modulation formats. Defaults to None.
    Returns:
        Tuple[chex.Array, chex.Array]: Transceiver noise array, Amplifier noise array
    """
    f = pathlib.Path(noise_data_filepath) if noise_data_filepath else (
            pathlib.Path(__file__).parents[1].absolute() / "data" / "gn_model" / "transceiver_amplifier_data.csv")
    noise_data = np.genfromtxt(f, delimiter=',')
    # Drop empty first row (headers) and column (name)
    noise_data = noise_data[1:, 1:]
    # Columns are: wavelength_min_nm,wavelength_max_nm,frequency_min_ghz,frequency_max_ghz,NF_ASE_dB,SNR_TRX_dB
    frequency_min_ghz = noise_data[:, 2]
    frequency_max_ghz = noise_data[:, 3]
    amplifier_noise_db = noise_data[:, 4]  # NF_ASE_dB
    transceiver_snr_db = noise_data[:, 5]  # SNR_TRX_dB

    # Define slot centres in GHz relative to central wavelength
    slot_centres = (jnp.arange(link_resources) - (link_resources - 1) / 2) * slot_size

    # Transform relative slot centres to absolute frequencies in GHz
    ref_frequency_ghz = c / ref_lambda / 1e9
    slot_frequencies_ghz = ref_frequency_ghz + slot_centres

    # Initialize output arrays
    transceiver_snr_array = jnp.zeros(link_resources)
    amplifier_noise_figure_array = jnp.zeros(link_resources)

    # For each slot, find which band it belongs to
    for i, freq in enumerate(slot_frequencies_ghz):
        # Find the band this frequency falls into
        for j in range(len(frequency_min_ghz)):
            if frequency_min_ghz[j] <= freq <= frequency_max_ghz[j]:
                transceiver_snr_array = transceiver_snr_array.at[i].set(transceiver_snr_db[j])
                amplifier_noise_figure_array = amplifier_noise_figure_array.at[i].set(amplifier_noise_db[j])
                break
        else:
            # If frequency is outside all bands, could raise error or use default
            raise ValueError(f"Frequency {freq} GHz is outside the defined bands")

    return transceiver_snr_array, amplifier_noise_figure_array

init_virtual_topology_patterns(pattern_names)

Initialise virtual topology patterns. First 3 digits comprise the "action counter": first index is num unique nodes, second index is total steps, final is remaining steps until completion of request. Remaining digits define the topology pattern, with 1 to indicate links and other positive integers are node indices.

Parameters:

Name	Type	Description	Default
pattern_names	list	List of virtual topology pattern names	required

Returns:

Type	Description
Array	chex.Array: Array of virtual topology patterns

Source code in xlron/environments/env_funcs.py

def init_virtual_topology_patterns(pattern_names: str) -> chex.Array:
    """Initialise virtual topology patterns.
    First 3 digits comprise the "action counter": first index is num unique nodes, second index is total steps,
    final is remaining steps until completion of request.
    Remaining digits define the topology pattern, with 1 to indicate links and other positive integers are node indices.

    Args:
        pattern_names (list): List of virtual topology pattern names

    Returns:
        chex.Array: Array of virtual topology patterns
    """
    patterns = []
    # TODO - Allow 2 node requests in VONE (check if any modifications necessary other than below)
    #if "2_bus" in pattern_names:
    #    patterns.append([2, 1, 1, 2, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0])
    if "3_bus" in pattern_names:
        patterns.append([3, 2, 2, 2, 1, 3, 1, 4])
    if "3_ring" in pattern_names:
        patterns.append([3, 3, 3, 2, 1, 3, 1, 4, 1, 2])
    if "4_bus" in pattern_names:
        patterns.append([4, 3, 3, 2, 1, 3, 1, 4, 1, 5])
    if "4_ring" in pattern_names:
        patterns.append([4, 4, 4, 2, 1, 3, 1, 4, 1, 5, 1, 2])
    if "5_bus" in pattern_names:
        patterns.append([5, 4, 4, 2, 1, 3, 1, 4, 1, 5, 1, 6])
    if "5_ring" in pattern_names:
        patterns.append([5, 5, 5, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 2])
    if "6_bus" in pattern_names:
        patterns.append([6, 5, 5, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 7])
    max_length = max([len(pattern) for pattern in patterns])
    # Pad patterns with zeroes to match longest
    for pattern in patterns:
        pattern.extend([0]*(max_length-len(pattern)))
    return jnp.array(patterns, dtype=SMALL_INT_DTYPE)

init_vone_request_array(params)

Initialize request array either with uniform resources

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_vone_request_array(params: EnvParams):
    """Initialize request array either with uniform resources"""
    return jnp.zeros((2, params.max_edges*2+1,), dtype=MED_INT_DTYPE)

make_graph(topology_name='conus', topology_directory=None)

Create graph from topology definition. Topologies must be defined in JSON format in the topologies directory and named as the topology name with .json extension.

Parameters:

Name	Type	Description	Default
topology_name	str	topology name	'conus'
topology_directory	str	topology directory	None

Returns:

Name	Type	Description
graph		graph

Source code in xlron/environments/env_funcs.py

def make_graph(topology_name: str = "conus", topology_directory: str = None):
    """Create graph from topology definition.
    Topologies must be defined in JSON format in the topologies directory and
    named as the topology name with .json extension.

    Args:
        topology_name: topology name
        topology_directory: topology directory

    Returns:
        graph: graph
    """
    topology_path = pathlib.Path(topology_directory) if topology_directory else (
            pathlib.Path(__file__).parents[1].absolute() / "data" / "topologies")
    # Create topology
    if topology_name == "4node":
        # 4 node ring
        graph = nx.from_numpy_array(np.array([[0, 1, 0, 1],
                                            [1, 0, 1, 0],
                                               [0, 1, 0, 1],
                                               [1, 0, 1, 0]]))
        # Add edge weights to graph
        nx.set_edge_attributes(graph, {(0, 1): 4, (1, 2): 3, (2, 3): 2, (3, 0): 1}, "weight")
    elif topology_name == "7node":
        # 7 node ring
        graph = nx.from_numpy_array(jnp.array([[0, 1, 0, 0, 0, 0, 1],
                                               [1, 0, 1, 0, 0, 0, 0],
                                               [0, 1, 0, 1, 0, 0, 0],
                                               [0, 0, 1, 0, 1, 0, 0],
                                               [0, 0, 0, 1, 0, 1, 0],
                                               [0, 0, 0, 0, 1, 0, 1],
                                               [1, 0, 0, 0, 0, 1, 0]]))
        # Add edge weights to graph
        nx.set_edge_attributes(graph, {(0, 1): 4, (1, 2): 3, (2, 3): 2, (3, 4): 1, (4, 5): 2, (5, 6): 3, (6, 0): 4}, "weight")
    else:
        with open(topology_path / f"{topology_name}.json") as f:
            graph = nx.node_link_graph(json.load(f))
    return graph

mask_nodes(state, num_nodes)

Returns mask of valid actions for node selection. 1 for valid action, 0 for invalid action.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
num_nodes	Scalar	Number of nodes	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_nodes(state: EnvState, num_nodes: chex.Scalar) -> EnvState:
    """Returns mask of valid actions for node selection. 1 for valid action, 0 for invalid action.

    Args:
        state: Environment state
        num_nodes: Number of nodes

    Returns:
        state: Updated environment state
    """
    total_actions = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.action_counter, 1, 1))
    remaining_actions = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.action_counter, 2, 1))
    full_request = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.request_array, 0, 1))
    virtual_topology = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.request_array, 1, 1))
    request = jax.lax.dynamic_slice_in_dim(full_request, (remaining_actions - 1) * 2, 3)
    node_request_s = jax.lax.dynamic_slice_in_dim(request, 2, 1)
    node_request_d = jax.lax.dynamic_slice_in_dim(request, 0, 1)
    prev_action = jax.lax.dynamic_slice_in_dim(state.action_history, (remaining_actions) * 2, 3)
    prev_dest = jax.lax.dynamic_slice_in_dim(prev_action, 0, 1)
    node_indices = jnp.arange(0, num_nodes)
    # Get requested indices from request array virtual topology
    requested_indices = jax.lax.dynamic_slice_in_dim(virtual_topology, (remaining_actions-1)*2, 3)
    requested_index_d = jax.lax.dynamic_slice_in_dim(requested_indices, 0, 1)
    # Get index of previous selected node
    prev_selected_node = jnp.where(virtual_topology == requested_index_d, state.action_history, jnp.full(virtual_topology.shape, -1))
    # will be current index if node only occurs once in virtual topology or will be different index if occurs more than once
    prev_selected_index = jnp.argmax(prev_selected_node).astype(MED_INT_DTYPE)
    prev_selected_node_d = jax.lax.dynamic_slice_in_dim(state.action_history, prev_selected_index, 1)

    # If first action, source and dest both to be assigned -> just mask all nodes based on resources
    # Thereafter, source must be previous dest. Dest can be any node (except previous allocations).
    state = state.replace(
        node_mask_s=jax.lax.cond(
            jnp.equal(remaining_actions, total_actions),
            lambda x: jnp.where(
                state.node_capacity_array >= node_request_s,
                x,
                jnp.zeros(num_nodes)
            ),
            lambda x: jnp.where(
                node_indices == prev_dest,
                x,
                jnp.zeros(num_nodes)
            ),
            jnp.ones(num_nodes),
        )
    )
    state = state.replace(
        node_mask_d=jnp.where(
            state.node_capacity_array >= node_request_d,
            jnp.ones(num_nodes),
            jnp.zeros(num_nodes)
        )
    )
    # If not first move, set node_mask_d to zero wherever node_mask_s is 1
    # to avoid same node selection for s and d
    state = state.replace(
        node_mask_d=jax.lax.cond(
            jnp.equal(remaining_actions, total_actions),
            lambda x: x,
            lambda x: jnp.where(
                state.node_mask_s == 1,
                jnp.zeros(num_nodes),
                x
            ),
            state.node_mask_d,
        )
    )

    def mask_previous_selections(i, val):
        # Disallow previously allocated nodes
        update_slice = lambda j, x: jax.lax.dynamic_update_slice_in_dim(x, jnp.array([0.]), j, axis=0)
        val = jax.lax.cond(
            i % 2 == 0,
            lambda x: update_slice(x[0][i], x[1]),  # i is node request index
            lambda x: update_slice(x[0][i+1], x[1]),  # i is slot request index (so add 1 to get next node)
            (state.action_history, val),
        )
        return val

    state = state.replace(
        node_mask_d=jax.lax.fori_loop(
            remaining_actions*2,
            state.action_history.shape[0]-1,
            mask_previous_selections,
            state.node_mask_d
        )
    )
    # If requested node index is new then disallow previously allocated nodes
    # If not new, then must match previously allocated node for that index
    state = state.replace(
        node_mask_d=jax.lax.cond(
            jnp.squeeze(prev_selected_node_d) >= 0,
            lambda x: jnp.where(
                node_indices == prev_selected_node_d,
                x[1],
                x[0],
            ),
            lambda x: x[2],
            (jnp.zeros(num_nodes), jnp.ones(num_nodes), state.node_mask_d),
        )
    )
    return state

mask_slots(state, params, request)

Returns binary mask of valid actions. 1 for valid action, 0 for invalid action.

1. Check request for source and destination nodes
2. For each path:
   • Get current slots on path (with padding on end to avoid out of bounds)
   • Get mask for required slots on path
   • Multiply through current slots with required slots mask to check if slots available on path
   • Remove padding from mask
   • Return path mask
3. Update total mask with path mask
4. If aggregate_slots > 1, aggregate slot mask to reduce action space

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_slots(state: EnvState, params: EnvParams, request: chex.Array) -> EnvState:
    """Returns binary mask of valid actions. 1 for valid action, 0 for invalid action.

    1. Check request for source and destination nodes
    2. For each path:
        - Get current slots on path (with padding on end to avoid out of bounds)
        - Get mask for required slots on path
        - Multiply through current slots with required slots mask to check if slots available on path
        - Remove padding from mask
        - Return path mask
    3. Update total mask with path mask
    4. If aggregate_slots > 1, aggregate slot mask to reduce action space

    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_mask = jnp.zeros((params.link_resources * params.k_paths), dtype=LARGE_FLOAT_DTYPE)

    def mask_path(i, mask):
        # Get slots for path
        slots = get_path_slots(state.link_slot_array, params, nodes_sd, i)
        # Add padding to slots at end
        slots = jnp.concatenate((slots, jnp.ones(params.max_slots, dtype=LARGE_FLOAT_DTYPE)))
        # Convert bandwidth to slots for each path
        spectral_efficiency = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else one
        requested_slots = required_slots(requested_datarate, spectral_efficiency, params.slot_size, guardband=params.guardband)
        # Get mask used to check if request will fit slots
        request_mask = get_request_mask(requested_slots[0], params)

        def check_slots_available(j, val):
            # Multiply through by request mask to check if slots available
            request_slice = jax.lax.dynamic_slice(val, (j,), (params.max_slots,))
            slot_sum = jnp.sum(request_mask * request_slice, promote_integers=False) <= zero
            slot_sum = slot_sum.reshape((1,)).astype(LARGE_FLOAT_DTYPE)
            return jax.lax.dynamic_update_slice(val, slot_sum, (j,))

        # Mask out slots that are not valid
        path_mask = jax.lax.fori_loop(
            0,
            int(params.link_resources+1),  # No need to check last requested_slots-1 slots
            check_slots_available,
            slots,
        )
        # Cut off padding
        path_mask = jax.lax.dynamic_slice(path_mask, (0,), (params.link_resources,))
        # Update total mask with path mask
        mask = jax.lax.dynamic_update_slice(mask, path_mask, (i * params.link_resources,))
        return mask

    # Loop over each path
    link_slot_mask = jax.lax.fori_loop(0, params.k_paths, mask_path, init_mask)
    if params.aggregate_slots > 1:
        # Full link slot mask is used in process_path_action to get the correct slot from the aggregated slot action
        state = state.replace(full_link_slot_mask=link_slot_mask)
        link_slot_mask, _ = aggregate_slots(link_slot_mask.reshape(params.k_paths, -1), params)
        link_slot_mask = link_slot_mask.reshape(-1)
    state = state.replace(link_slot_mask=link_slot_mask)
    return state

mask_slots_rmsa_gn_model(state, params, request)

For use in RSAGNModelEnv. 1. For each path: 1.1 Get path slots 1.2 Get launch power

Parameters:

Name	Type	Description	Default
state	RSAGNModelEnvState	Environment state	required
params	RSAGNModelEnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_slots_rmsa_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams, request: chex.Array) -> EnvState:
    """For use in RSAGNModelEnv.
    1. For each path:
        1.1 Get path slots
        1.2 Get launch power


    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_mask = jnp.zeros((params.link_resources * params.k_paths))

    def mask_path(i, mask):
        path = get_paths(params, nodes_sd)[i]
        # Get slots for path
        slots = get_path_slots(state.link_slot_array, params, nodes_sd, i)
        # Add padding to slots at end
        # 0 means slot is free, 1 is occupied
        slots = jnp.concatenate((slots, jnp.ones(params.max_slots)))
        launch_power = get_launch_power(state, i, state.launch_power_array[i], params)
        lightpath_index = get_lightpath_index(params, nodes_sd, i)

        # This function checks through each available modulation format, checks the first and last available slots,
        # calculates the SNR, checks it meets the requirements, and returns the resulting mask
        def check_modulation_format(mod_format_index, init_path_mask):
            se = params.modulations_array.val[mod_format_index][1]
            req_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)[0]
            bandwidth_per_subchannel = params.slot_size
            req_snr = params.modulations_array.val[mod_format_index][2] + params.snr_margin
            # Get mask used to check if request will fit slots
            request_mask = get_request_mask(req_slots, params)

            def check_slots_available(j, val):
                # Multiply through by request mask to check if slots available
                slot_sum = jnp.sum(request_mask * jax.lax.dynamic_slice(val, (j,), (params.max_slots,)), promote_integers=False) <= 0
                slot_sum = slot_sum.reshape((1,)).astype(LARGE_FLOAT_DTYPE)
                return jax.lax.dynamic_update_slice(val, slot_sum, (j,))

            # Mask out slots that are not valid
            slot_mask = jax.lax.fori_loop(
                0,
                int(params.link_resources + 1),  # No need to check last requested_slots-1 slots
                check_slots_available,
                slots,
            )
            # Cut off padding
            slot_mask = jax.lax.dynamic_slice(slot_mask, (0,), (params.link_resources,))
            # Check first and last available slots for suitability
            ff_path_mask = jnp.concatenate((slot_mask, jnp.ones((1,))), axis=0)
            lf_path_mask = jnp.concatenate((jnp.ones((1,)), slot_mask), axis=0)
            first_available_slot_index = jnp.argmax(ff_path_mask)
            last_available_slot_index = params.link_resources - jnp.argmax(jnp.flip(lf_path_mask)) - 1
            # Assign "req_slots" subchannels (each with bandwidth = slot width) for the first and last possible slots
            ff_temp_state = state.replace(
                channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, first_available_slot_index, req_slots, bandwidth_per_subchannel),
                channel_power_array=vmap_set_path_links(state.channel_power_array, path, first_available_slot_index, req_slots, launch_power),
                path_index_array=vmap_set_path_links(state.path_index_array, path, first_available_slot_index, req_slots, lightpath_index),
                modulation_format_index_array=vmap_set_path_links(state.modulation_format_index_array, path, first_available_slot_index, req_slots, mod_format_index),
            )
            lf_temp_state = state.replace(
                channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, last_available_slot_index, req_slots, bandwidth_per_subchannel),
                channel_power_array=vmap_set_path_links(state.channel_power_array, path, last_available_slot_index, req_slots, launch_power),
                path_index_array=vmap_set_path_links(state.path_index_array, path, last_available_slot_index, req_slots, lightpath_index),
                modulation_format_index_array=vmap_set_path_links(state.modulation_format_index_array, path, last_available_slot_index, req_slots, mod_format_index),
            )
            ff_temp_state = ff_temp_state.replace(link_snr_array=get_snr_link_array(ff_temp_state, params))
            lf_temp_state = lf_temp_state.replace(link_snr_array=get_snr_link_array(lf_temp_state, params))
            # Take the minimum value of SNR from all the subchannels
            ff_snr_value = get_minimum_snr_of_channels_on_path(
                ff_temp_state, path, first_available_slot_index, req_slots, params
            )
            lf_snr_value = get_minimum_snr_of_channels_on_path(
                lf_temp_state, path, last_available_slot_index, req_slots, params
            )
            # Check that other paths SNR is still sufficient (True if failure)
            ff_snr_check = 1 - check_action_rmsa_gn_model(ff_temp_state, None, params)
            lf_snr_check = 1 - check_action_rmsa_gn_model(lf_temp_state, None, params)
            ff_check = (ff_snr_value >= req_snr) * ff_snr_check
            lf_check = (lf_snr_value >= req_snr) * lf_snr_check

            slot_indices = jnp.arange(params.link_resources, dtype=MED_INT_DTYPE)
            mod_format_mask = jnp.where(slot_indices == first_available_slot_index, ff_check, False)
            mod_format_mask = jnp.where(slot_indices == last_available_slot_index, lf_check, mod_format_mask)
            path_mask = jnp.where(mod_format_mask, mod_format_index, init_path_mask)
            # jax.debug.print("ff_snr_check {}", ff_snr_check, ordered=True)
            # jax.debug.print("lf_snr_check {}", lf_snr_check, ordered=True)
            # jax.debug.print("ff_snr_value {}", ff_snr_value, ordered=True)
            # jax.debug.print("lf_snr_value {}", lf_snr_value, ordered=True)
            # jax.debug.print("first_available_slot_index {}", first_available_slot_index, ordered=True)
            # jax.debug.print("last_available_slot_index {}", last_available_slot_index, ordered=True)
            # jax.debug.print("req_snr {}", req_snr, ordered=True)
            # jax.debug.print("mod_format_mask {}", mod_format_mask, ordered=True)
            # jax.debug.print("path_mask {}", path_mask, ordered=True)
            return path_mask

        path_mask = jax.lax.fori_loop(0, params.modulations_array.val.shape[0], check_modulation_format, jnp.full((params.link_resources,), -1., dtype=LARGE_FLOAT_DTYPE))

        # Update total mask with path mask
        mask = jax.lax.dynamic_update_slice(mask, path_mask, (i * params.link_resources,))
        return mask

    # Loop over each path
    mod_format_mask = jax.lax.fori_loop(0, params.k_paths, mask_path, init_mask)
    link_slot_mask = jnp.where(mod_format_mask >= 0, 1.0, 0.0)
    if params.aggregate_slots > 1:
        # Full link slot mask is used in process_path_action to get the correct slot from the aggregated slot action
        state = state.replace(full_link_slot_mask=link_slot_mask)
        link_slot_mask, _ = aggregate_slots(link_slot_mask.reshape(params.k_paths, -1), params)
        link_slot_mask = link_slot_mask.reshape(-1)
    state = state.replace(
        link_slot_mask=link_slot_mask,
        mod_format_mask=mod_format_mask,
    )
    return state

mask_slots_rwalr(state, params, request)

For use in RWALightpathReuseEnv. Each lightpath has a maximum capacity defined in path_capacity_array. This is updated when a lightpath is assigned. If remaining path capacity is less than current request, corresponding link-slots are masked out. If link-slot is in use by another lightpath for a different source and destination node (even if not full) it is masked out. Step 1: - Mask out slots that are not valid based on path capacity (check link_capacity_array) Step 2: - Mask out slots that are not valid based on lightpath reuse (check path_index_array)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_slots_rwalr(state: EnvState, params: EnvParams, request: chex.Array) -> EnvState:
    """For use in RWALightpathReuseEnv.
    Each lightpath has a maximum capacity defined in path_capacity_array. This is updated when a lightpath is assigned.
    If remaining path capacity is less than current request, corresponding link-slots are masked out.
    If link-slot is in use by another lightpath for a different source and destination node (even if not full) it is masked out.
    Step 1:
    - Mask out slots that are not valid based on path capacity (check link_capacity_array)
    Step 2:
    - Mask out slots that are not valid based on lightpath reuse (check path_index_array)

    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_mask = jnp.zeros((params.link_resources * params.k_paths))
    source, dest = nodes_sd
    path_start_index = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
    #jax.debug.print("path_start_index {}", path_start_index, ordered=True)
    #jax.debug.print("link_capacity_array {}", state.link_capacity_array, ordered=True)

    def mask_path(i, mask):
        # Step 1 - mask capacity
        capacity_mask = jnp.where(state.link_capacity_array < requested_datarate, 1., 0.)
        #jax.debug.print("capacity_mask {}", capacity_mask, ordered=True)
        capacity_slots = get_path_slots(capacity_mask, params, nodes_sd, i)
        #jax.debug.print("capacity_slots {}", capacity_slots, ordered=True)
        # Step 2 - mask lightpath reuse
        lightpath_index = path_start_index + i
        #jax.debug.print("lightpath_index {}", lightpath_index, ordered=True)
        lightpath_mask = jnp.where(state.path_index_array == lightpath_index, 0., 1.)  # Allow current lightpath
        #jax.debug.print("lightpath_mask {}", lightpath_mask, ordered=True)
        lightpath_mask = jnp.where(state.path_index_array == -1, 0., lightpath_mask)  # Allow empty slots
        #jax.debug.print("lightpath_mask {}", lightpath_mask, ordered=True)
        lightpath_slots = get_path_slots(lightpath_mask, params, nodes_sd, i)
        #jax.debug.print("lightpath_slots {}", lightpath_slots, ordered=True)
        # Step 3 combine masks
        path_mask = jnp.max(jnp.stack((capacity_slots, lightpath_slots)), axis=0)
        # Swap zeros for ones
        path_mask = jnp.where(path_mask == 0, 1., 0.)
        #jax.debug.print("path_mask {}", path_mask, ordered=True)
        mask = jax.lax.dynamic_update_slice(mask, path_mask, (i * params.link_resources,))
        return mask

    # Loop over each path
    link_slot_mask = jax.lax.fori_loop(0, params.k_paths, mask_path, init_mask)
    if params.aggregate_slots > 1:
        # Full link slot mask is used in process_path_action to get the correct slot from the aggregated slot action
        state = state.replace(full_link_slot_mask=link_slot_mask)
        link_slot_mask, _ = aggregate_slots(link_slot_mask.reshape(params.k_paths, -1), params)
        link_slot_mask = link_slot_mask.reshape(-1)
    state = state.replace(link_slot_mask=link_slot_mask)
    return state

normalise_traffic_matrix(traffic_matrix)

Normalise traffic matrix to sum to 1

Source code in xlron/environments/env_funcs.py

def normalise_traffic_matrix(traffic_matrix):
    """Normalise traffic matrix to sum to 1"""
    traffic_matrix /= jnp.sum(traffic_matrix, promote_integers=False)
    return traffic_matrix

pad_array(array, fill_value)

Pad a ragged multidimensional array to rectangular shape. Used for training on multiple topologies. Source: https://codereview.stackexchange.com/questions/222623/pad-a-ragged-multidimensional-array-to-rectangular-shape

Parameters:

Name	Type	Description	Default
array		array to pad	required
fill_value		value to fill with	required

Returns:

Name	Type	Description
result		padded array

Source code in xlron/environments/env_funcs.py

def pad_array(array, fill_value):
    """
    Pad a ragged multidimensional array to rectangular shape.
    Used for training on multiple topologies.
    Source: https://codereview.stackexchange.com/questions/222623/pad-a-ragged-multidimensional-array-to-rectangular-shape

    Args:
        array: array to pad
        fill_value: value to fill with

    Returns:
        result: padded array
    """

    def get_dimensions(array, level=0):
        yield level, len(array)
        try:
            for row in array:
                yield from get_dimensions(row, level + 1)
        except TypeError: #not an iterable
            pass

    def get_max_shape(array):
        dimensions = defaultdict(int)
        for level, length in get_dimensions(array):
            dimensions[level] = max(dimensions[level], length)
        return [value for _, value in sorted(dimensions.items())]

    def iterate_nested_array(array, index=()):
        try:
            for idx, row in enumerate(array):
                yield from iterate_nested_array(row, (*index, idx))
        except TypeError: # final level
            yield (*index, slice(len(array))), array

    dimensions = get_max_shape(array)
    result = np.full(dimensions, fill_value)
    for index, value in iterate_nested_array(array):
        result[index] = value
    return result

path_action_only(topology_pattern, action_counter, remaining_actions)

This is to check if node has already been assigned, therefore just need to assign slots (n=0)

Parameters:

Name	Type	Description	Default
topology_pattern	Array	Topology pattern	required
action_counter	Array	Action counter	required
remaining_actions	Scalar	Remaining actions	required

Returns:

Name	Type	Description
bool		True if only path action, False if node action

Source code in xlron/environments/env_funcs.py

def path_action_only(topology_pattern: chex.Array, action_counter: chex.Array, remaining_actions: chex.Scalar):
    """This is to check if node has already been assigned, therefore just need to assign slots (n=0)

    Args:
        topology_pattern: Topology pattern
        action_counter: Action counter
        remaining_actions: Remaining actions

    Returns:
        bool: True if only path action, False if node action
    """
    # Get topology segment to be assigned e.g. [2,1,4]
    topology_segment = jax.lax.dynamic_slice(topology_pattern, ((remaining_actions-1)*2, ), (3, ))
    topology_indices = jnp.arange(topology_pattern.shape[0])
    # Check if the latest node in the segment is found in "prev_assigned_topology"
    new_node_to_be_assigned = topology_segment[0]
    prev_assigned_topology = jnp.where(topology_indices > (action_counter[-1]-1)*2, topology_pattern, 0)
    nodes_already_assigned_check = jnp.any(jnp.sum(jnp.where(prev_assigned_topology == new_node_to_be_assigned, 1, 0)) > 0)
    return nodes_already_assigned_check

poisson(key, lam, shape=(), dtype=dtypes.float_)

Sample Exponential random values with given shape and float dtype.

The values are distributed according to the probability density function:

.. math:: f(x) = \lambda e^{-\lambda x}

on the domain :math:0 \le x < \infty.

Args: key: a PRNG key used as the random key. lam: a positive float32 or float64 Tensor indicating the rate parameter shape: optional, a tuple of nonnegative integers representing the result shape. Default (). dtype: optional, a float dtype for the returned values (default float64 if jax_enable_x64 is true, otherwise float32).

Returns: A random array with the specified shape and dtype.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 2, 3))
def poisson(key: Union[Array, prng.PRNGKeyArray],
            lam: ArrayLike,
            shape: Shape = (),
            dtype: DTypeLike = dtypes.float_) -> Array:
    r"""Sample Exponential random values with given shape and float dtype.

    The values are distributed according to the probability density function:

    .. math::
     f(x) = \lambda e^{-\lambda x}

    on the domain :math:`0 \le x < \infty`.

    Args:
    key: a PRNG key used as the random key.
    lam: a positive float32 or float64 `Tensor` indicating the rate parameter
    shape: optional, a tuple of nonnegative integers representing the result
      shape. Default ().
    dtype: optional, a float dtype for the returned values (default float64 if
      jax_enable_x64 is true, otherwise float32).

    Returns:
    A random array with the specified shape and dtype.
    """
    key, _ = jax._src.random._check_prng_key(key)
    if not dtypes.issubdtype(dtype, np.floating):
        raise ValueError(f"dtype argument to `exponential` must be a float "
                       f"dtype, got {dtype}")
    dtype = dtypes.canonicalize_dtype(dtype)
    shape = core.canonicalize_shape(shape)
    return _poisson(key, lam, shape, dtype)

process_path_action(state, params, path_action)

Process path action to get path index and initial slot index.

Parameters:

Name	Type	Description	Default
state	State	current state	required
params	Params	environment parameters	required
path_action	int	path action	required

Returns:

Name	Type	Description
int	(Array, Array)	path index
int	(Array, Array)	initial slot index

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def process_path_action(state: EnvState, params: EnvParams, path_action: chex.Array) -> (chex.Array, chex.Array):
    """Process path action to get path index and initial slot index.

    Args:
        state (State): current state
        params (Params): environment parameters
        path_action (int): path action

    Returns:
        int: path index
        int: initial slot index
    """
    num_slot_actions = math.ceil(params.link_resources/params.aggregate_slots)
    path_index = jnp.floor(path_action / num_slot_actions).astype(LARGE_INT_DTYPE).reshape(1)
    initial_aggregated_slot_index = jnp.mod(path_action, num_slot_actions).reshape(1)
    initial_slot_index = initial_aggregated_slot_index*params.aggregate_slots
    if params.aggregate_slots > 1:
        # Get the path mask do a dynamic slice and get the index of first unoccupied slot in the slice
        path_mask = jax.lax.dynamic_slice(state.full_link_slot_mask, path_index*params.link_resources, (params.link_resources,))
        path_mask_slice = jax.lax.dynamic_slice(path_mask, initial_slot_index, (params.aggregate_slots,))
        # Use argmax to get index of first 1 in slice of mask
        initial_slot_index = initial_slot_index + jnp.argmax(path_mask_slice).astype(MED_INT_DTYPE)
    return path_index[0], initial_slot_index[0]

read_rsa_request(request_array)

Read RSA request from request array. Return source-destination nodes and bandwidth request.

Parameters:

Name	Type	Description	Default
request_array	Array	request array	required

Returns:

Type	Description
Tuple[Array, Array]	Tuple[chex.Array, chex.Array]: source-destination nodes and bandwidth request

Source code in xlron/environments/env_funcs.py

def read_rsa_request(request_array: chex.Array) -> Tuple[chex.Array, chex.Array]:
    """Read RSA request from request array. Return source-destination nodes and bandwidth request.

    Args:
        request_array: request array

    Returns:
        Tuple[chex.Array, chex.Array]: source-destination nodes and bandwidth request
    """
    node_s = jax.lax.dynamic_slice(request_array, (0,), (1,))
    requested_datarate = jax.lax.dynamic_slice(request_array, (1,), (1,))
    node_d = jax.lax.dynamic_slice(request_array, (2,), (1,))
    nodes_sd = jnp.concatenate((node_s, node_d))
    return nodes_sd, requested_datarate

remove_expired_node_requests(state, params)

Check for values in node_departure_array that are less than the current time but greater than zero (negative time indicates the request is not yet finalised). If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase node_capacity_array by expired resources on each node.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_node_requests(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """Check for values in node_departure_array that are less than the current time but greater than zero
    (negative time indicates the request is not yet finalised).
    If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase
    node_capacity_array by expired resources on each node.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    mask = jnp.where(state.node_departure_array < jnp.squeeze(state.current_time), 1, 0)
    mask = jnp.where(0 < state.node_departure_array, mask, 0)
    expired_resources = jnp.sum(jnp.where(mask == 1, state.node_resource_array, 0), axis=1, promote_integers=False)
    state = state.replace(
        node_capacity_array=state.node_capacity_array + expired_resources,
        node_resource_array=jnp.where(mask == 1, 0, state.node_resource_array),
        node_departure_array=jnp.where(mask == 1, jnp.inf, state.node_departure_array)
    )
    return state

remove_expired_services_rmsa_gn_model(state, params)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rmsa_gn_model(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                      state.link_slot_departure_array - jnp.squeeze(current_time),
                                                      updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=jnp.where(mask_remove == one, zero, state.link_slot_array),
        link_slot_departure_array=updated_link_slot_departure_array,
        link_snr_array=jnp.where(mask_remove == one, zero, state.link_snr_array),
        path_index_array=jnp.where(mask_remove == one, -one, state.path_index_array),
        channel_centre_bw_array=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array),
        channel_power_array=jnp.where(mask_remove == one, zero, state.channel_power_array),
        modulation_format_index_array=jnp.where(mask_remove == one, -one, state.modulation_format_index_array),
        path_index_array_prev=jnp.where(mask_remove == one, -one, state.path_index_array_prev),
        channel_centre_bw_array_prev=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array_prev),
        channel_power_array_prev=jnp.where(mask_remove == one, zero, state.channel_power_array_prev),
        modulation_format_index_array_prev=jnp.where(mask_remove == one, -one, state.modulation_format_index_array_prev),
    )
    return state

remove_expired_services_rsa(state, params)

Check for values in link_slot_departure_array that are less than the current time but greater than zero (negative time indicates the request is not yet finalised). If found, set to zero in link_slot_array and link_slot_departure_array.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rsa(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """Check for values in link_slot_departure_array that are less than the current time but greater than zero
    (negative time indicates the request is not yet finalised).
    If found, set to zero in link_slot_array and link_slot_departure_array.

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_array = jnp.where(mask_remove == one, zero, state.link_slot_array)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                     state.link_slot_departure_array - jnp.squeeze(current_time),
                                                     updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=updated_link_slot_array,
        link_slot_departure_array=updated_link_slot_departure_array,
    )
    return state

remove_expired_services_rsa_gn_model(state, params)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rsa_gn_model(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                      state.link_slot_departure_array - jnp.squeeze(current_time),
                                                      updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=jnp.where(mask_remove == one, zero, state.link_slot_array),
        link_slot_departure_array=updated_link_slot_departure_array,
        link_snr_array=jnp.where(mask_remove == one, zero, state.link_snr_array),
        path_index_array=jnp.where(mask_remove == one, -one, state.path_index_array),
        channel_centre_bw_array=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array),
        channel_power_array=jnp.where(mask_remove == one, zero, state.channel_power_array),
        path_index_array_prev=jnp.where(mask_remove == one, -one, state.path_index_array_prev),
        channel_centre_bw_array_prev=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array_prev),
        channel_power_array_prev=jnp.where(mask_remove == one, zero, state.channel_power_array_prev),
    )
    if params.monitor_active_lightpaths:
        # The active_lightpaths_array is set to -1 when the lightpath is not active
        # The active_lightpaths_array_departure is set to 0 when the lightpath is not active
        # (active_lightpaths_array is used to calculate the total throughput)
        mask_remove = jnp.where(
            (zero <= state.active_lightpaths_array_departure) & (state.active_lightpaths_array_departure <= jnp.squeeze(current_time)),
            one, zero)
        state = state.replace(
            active_lightpaths_array=jnp.where(mask_remove == one, -one, state.active_lightpaths_array),
            active_lightpaths_array_departure=jnp.where(mask_remove == one, zero, state.active_lightpaths_array_departure),
        )
    return state

remove_expired_services_rwalr(state, params)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rwalr(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                     state.link_slot_departure_array - jnp.squeeze(current_time),
                                                     updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=jnp.where(mask_remove == one, zero, state.link_slot_array),
        path_index_array=jnp.where(mask_remove == one, -one, state.path_index_array),
        link_slot_departure_array=updated_link_slot_departure_array,
    )
    return state

required_slots(bitrate, se, channel_width, guardband=1)

Calculate required slots for a given bitrate and spectral efficiency.

Parameters:

Name	Type	Description	Default
bit_rate	float	Bit rate in Gbps	required
se	float	Spectral efficiency in bps/Hz	required
channel_width	float	Channel width in GHz	required
guardband	int	Guard band. Defaults to 1.	1

Returns:

Name	Type	Description
int	Array	Required slots

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2, 3))
def required_slots(bitrate: float, se: int, channel_width: float, guardband: int = 1) -> chex.Array:
    """Calculate required slots for a given bitrate and spectral efficiency.

    Args:
        bit_rate (float): Bit rate in Gbps
        se (float): Spectral efficiency in bps/Hz
        channel_width (float): Channel width in GHz
        guardband (int, optional): Guard band. Defaults to 1.

    Returns:
        int: Required slots
    """
    # If bitrate is zero, then required slots should be zero
    return ((jnp.ceil(bitrate/(se*channel_width))+guardband) * (one - (bitrate == zero))).astype(MED_INT_DTYPE)

set_band_gaps(link_slot_array, params, val)

Set band gaps in link slot array Args: link_slot_array (chex.Array): Link slot array params (RSAGNModelEnvParams): Environment parameters val (int): Value to set Returns: chex.Array: Link slot array with band gaps

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 2))
def set_band_gaps(link_slot_array: chex.Array, params: RSAGNModelEnvParams, val: int) -> chex.Array:
    """Set band gaps in link slot array
    Args:
        link_slot_array (chex.Array): Link slot array
        params (RSAGNModelEnvParams): Environment parameters
        val (int): Value to set
    Returns:
        chex.Array: Link slot array with band gaps
    """
    # Create array that is size of link_slot array with values of column index
    mask = jnp.arange(params.link_resources)
    mask = jnp.tile(mask, (params.num_links, 1))
    def set_band_gap(i, arr):
        gap_start = params.gap_starts.val[i]
        gap_end = gap_start + params.gap_widths.val[i]
        condition = jnp.logical_and(arr >= gap_start, arr < gap_end)
        arr = jnp.where(condition, -one, arr)
        return arr
    mask = jax.lax.fori_loop(0, params.gap_widths.val.shape[0], set_band_gap, mask)
    link_slot_array = jnp.where(mask == -one, val, link_slot_array)
    return link_slot_array

undo_action_rmsa_gn_model(state, params)

Undo action for RMSA GN model Args: state (EnvState): Environment state action (chex.Array): Action array params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def undo_action_rmsa_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> EnvState:
    """Undo action for RMSA GN model
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action array
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    state = undo_action_rsa(state, params)  # Undo link_slot_array and link_slot_departure_array
    state = state.replace(
        link_slot_array=set_band_gaps(state.link_slot_array, params, -one),  # Set C+L band gap
        channel_centre_bw_array=state.channel_centre_bw_array_prev,
        path_index_array=state.path_index_array_prev,
        channel_power_array=state.channel_power_array_prev,
        modulation_format_index_array=state.modulation_format_index_array_prev,
    )
    return state

undo_action_rsa(state, params)

If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation. Partial resource allocation is indicated by negative time in link slot departure array. Check for values in link_slot_departure_array that are less than zero. If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def undo_action_rsa(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation.
    Partial resource allocation is indicated by negative time in link slot departure array.
    Check for values in link_slot_departure_array that are less than zero.
    If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    # If departure array is negative, then undo the action
    mask = jnp.where(state.link_slot_departure_array < zero, one, zero)
    # If link slot array is < -1, then undo the action
    # (departure might be positive because existing service had holding time after current)
    # e.g. (time_in_array = t1 - t2) where t2 < t1 and t2 = current_time + holding_time
    # but link_slot_array = -2 due to double allocation, so undo the action
    mask = jnp.where(state.link_slot_array < -one, one, mask)
    state = state.replace(
        link_slot_array=jnp.where(mask == one, state.link_slot_array + one, state.link_slot_array),
        link_slot_departure_array=jnp.where(
            mask == one,
            state.link_slot_departure_array + state.current_time + state.holding_time,
            state.link_slot_departure_array),
        total_bitrate=state.total_bitrate + read_rsa_request(state.request_array)[1][0],
    )
    return state

undo_action_rsa_gn_model(state, params)

Undo action for RSA GN model Args: state (EnvState): Environment state action (chex.Array): Action array params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def undo_action_rsa_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> EnvState:
    """Undo action for RSA GN model
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action array
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    state = undo_action_rsa(state, params)  # Undo link_slot_array and link_slot_departure_array
    state = state.replace(
        link_slot_array=set_band_gaps(state.link_slot_array, params, -one),  # Set C+L band gap
        channel_centre_bw_array=state.channel_centre_bw_array_prev,
        path_index_array=state.path_index_array_prev,
        channel_power_array=state.channel_power_array_prev,
    )
    if params.monitor_active_lightpaths:
        # If departure array is negative, then undo the action
        mask = jnp.where(state.active_lightpaths_array_departure < zero, one,  zero)
        state = state.replace(
            active_lightpaths_array=jnp.where(mask == one, -one, state.active_lightpaths_array),
            active_lightpaths_array_departure=jnp.where(
                mask == one,
                state.active_lightpaths_array_departure + state.current_time + state.holding_time,
                state.active_lightpaths_array_departure),
        )
    return state

undo_action_rwalr(state, params)

If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation. Partial resource allocation is indicated by negative time in link slot departure array. Check for values in link_slot_departure_array that are less than zero. If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def undo_action_rwalr(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation.
    Partial resource allocation is indicated by negative time in link slot departure array.
    Check for values in link_slot_departure_array that are less than zero.
    If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    # If departure array is negative, then undo the action
    mask = jnp.where(state.link_slot_departure_array < zero, one, zero)
    # If link slot array is < -1, then undo the action
    # (departure might be positive because existing service had holding time after current)
    # e.g. (time_in_array = t1 - t2) where t2 < t1 and t2 = current_time + holding_time
    # but link_slot_array = -2 due to double allocation, so undo the action
    mask = jnp.where(state.link_slot_array < -one, one, mask)
    state = state.replace(
        link_slot_array=jnp.where(mask == one, state.link_slot_array + one, state.link_slot_array),
        link_slot_departure_array=jnp.where(
            mask == one,
            state.link_slot_departure_array + state.current_time + state.holding_time,
            state.link_slot_departure_array),
        total_bitrate=state.total_bitrate + read_rsa_request(state.request_array)[1][0]
    )
    return state

undo_node_action(state)

If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation. Partial resource allocation is indicated by negative time in node departure array. Check for values in node_departure_array that are less than zero. If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase node_capacity_array by expired resources on each node.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def undo_node_action(state: EnvState) -> EnvState:
    """If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation.
    Partial resource allocation is indicated by negative time in node departure array.
    Check for values in node_departure_array that are less than zero.
    If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase
    node_capacity_array by expired resources on each node.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    # TODO - Check that node resource clash doesn't happen (so time is always negative after implementation)
    #  and undoing always succeeds with negative time
    mask = jnp.where(state.node_departure_array < 0, 1, 0)
    resources = jnp.sum(jnp.where(mask == 1, state.node_resource_array, 0), axis=1, promote_integers=False)
    state = state.replace(
        node_capacity_array=state.node_capacity_array + resources,
        node_resource_array=jnp.where(mask == 1, 0, state.node_resource_array),
        node_departure_array=jnp.where(mask == 1, jnp.inf, state.node_departure_array),
    )
    return state

update_action_history(action_history, action_counter, action)

Update action history by adding action to first available index starting from the end.

Parameters:

Name	Type	Description	Default
action_history	Array	Action history	required
action_counter	Array	Action counter	required
action	Array	Action to add to history	required

Returns:

Type	Description
Array	Updated action_history

Source code in xlron/environments/env_funcs.py

def update_action_history(action_history: chex.Array, action_counter: chex.Array, action: chex.Array) -> chex.Array:
    """Update action history by adding action to first available index starting from the end.

    Args:
        action_history: Action history
        action_counter: Action counter
        action: Action to add to history

    Returns:
        Updated action_history
    """
    return jax.lax.dynamic_update_slice(action_history, jnp.flip(action, axis=0).astype(MED_INT_DTYPE), ((action_counter[-1]-1)*2,))

update_active_lightpaths_array(state, path_index, initial_slot_index, num_slots)

Update active lightpaths array with new path index. Find the first index of the array with value -1 and replace with path index. Args: state (RSAGNModelEnvState): Environment state path_index (int): Path index to add to active lightpaths array Returns: jnp.array: Updated active lightpaths array

Source code in xlron/environments/env_funcs.py

def update_active_lightpaths_array(state: RSAGNModelEnvState, path_index: int, initial_slot_index: int, num_slots: int) -> chex.Array:
    """Update active lightpaths array with new path index.
    Find the first index of the array with value -1 and replace with path index.
    Args:
        state (RSAGNModelEnvState): Environment state
        path_index (int): Path index to add to active lightpaths array
    Returns:
        jnp.array: Updated active lightpaths array
    """
    first_empty_index = jnp.argmin(state.active_lightpaths_array[:, 0])  # Just look at the first column
    return jax.lax.dynamic_update_slice(state.active_lightpaths_array, jnp.array([[path_index, initial_slot_index, num_slots[0]]]), (first_empty_index, 0))

update_active_lightpaths_array_departure(state, time)

Update active lightpaths array with new path index. Find the first index of the array with value -1 and replace with path index. Args: state (RSAGNModelEnvState): Environment state time (float): Departure time Returns: jnp.array: Updated active lightpaths array

Source code in xlron/environments/env_funcs.py

def update_active_lightpaths_array_departure(state: RSAGNModelEnvState, time: float) -> chex.Array:
    """Update active lightpaths array with new path index.
    Find the first index of the array with value -1 and replace with path index.
    Args:
        state (RSAGNModelEnvState): Environment state
        time (float): Departure time
    Returns:
        jnp.array: Updated active lightpaths array
    """
    first_empty_index = jnp.argmin(state.active_lightpaths_array[:, 0])  # Just look at the first column
    return jax.lax.dynamic_update_slice(state.active_lightpaths_array_departure, jnp.stack((time, time, time)), (first_empty_index, 0))

update_graph_tuple(state, params)

Update graph tuple for use with Jraph GNNs. Edge and node features are updated from link_slot_array and node_capacity_array respectively. Global features are updated as request_array. Args: state (EnvState): Environment state params (EnvParams): Environment parameters Returns: state (EnvState): Environment state with updated graph tuple

Source code in xlron/environments/env_funcs.py

def update_graph_tuple(state: EnvState, params: EnvParams):
    """Update graph tuple for use with Jraph GNNs.
    Edge and node features are updated from link_slot_array and node_capacity_array respectively.
    Global features are updated as request_array.
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
    Returns:
        state (EnvState): Environment state with updated graph tuple
    """
    # Get source and dest from request array
    source_dest, datarate = read_rsa_request(state.request_array)
    source, dest = source_dest[0], source_dest[2]
    # Global feature is normalised data rate of current request
    globals = jnp.array([datarate / jnp.max(params.values_bw.val)], dtype=LARGE_FLOAT_DTYPE)
    # One-hot encode source and destination
    source_dest_features = jnp.zeros((params.num_nodes, 2), dtype=LARGE_FLOAT_DTYPE)
    source_dest_features = source_dest_features.at[source.astype(MED_INT_DTYPE), 0].set(1)
    source_dest_features = source_dest_features.at[dest.astype(MED_INT_DTYPE), 1].set(-1)
    spectral_features = state.graph.nodes[..., :3]
    holding_time_edge_features = state.link_slot_departure_array / params.mean_service_holding_time

    if params.__class__.__name__ in ["RSAGNModelEnvParams", "RMSAGNModelEnvParams"]:
        # Normalize by max parameters (converted to linear units)
        max_power = isrs_gn_model.from_dbm(params.max_power)
        normalized_power = jnp.round(state.channel_power_array / max_power, 3)
        max_snr = isrs_gn_model.from_db(params.max_snr)
        normalized_snr = jnp.round(state.link_snr_array / max_snr, 3)
        edge_features = jnp.stack([normalized_snr, normalized_power], axis=-1)
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)
    elif params.__class__.__name__ == "VONEEnvParams":
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        node_features = getattr(state, "node_capacity_array", jnp.zeros(params.num_nodes))
        node_features = node_features.reshape(-1, 1)
        node_features = jnp.concatenate([node_features, spectral_features, source_dest_features], axis=-1)
    else:
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)

    if params.disable_node_features:
        node_features = jnp.zeros((1,), dtype=LARGE_FLOAT_DTYPE)

    edge_features = edge_features if params.directed_graph else jnp.repeat(edge_features, 2, axis=0)
    graph = state.graph._replace(nodes=node_features, edges=edge_features, globals=globals)
    state = state.replace(graph=graph)
    return state

update_node_array(node_indices, array, node, request)

Used to udated selected_nodes array with new requested resources on each node, for use in

Source code in xlron/environments/env_funcs.py

def update_node_array(node_indices, array, node, request):
    """Used to udated selected_nodes array with new requested resources on each node, for use in """
    return jnp.where(node_indices == node, array-request, array)

vmap_set_path_links(link_slot_array, path, initial_slot, num_slots, value)

Set relevant slots along links in path to val.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
path	Array	Path (row from path-link array that indicates links used by path)	required
initial_slot	int	Initial slot	required
num_slots	int	Number of slots	required
value	int	Value to set on link slot array	required

Returns:

Type	Description
Array	Updated link slot array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_set_path_links(link_slot_array: chex.Array, path: chex.Array, initial_slot: int, num_slots: int, value: int) -> chex.Array:
    """Set relevant slots along links in path to val.

    Args:
        link_slot_array: Link slot array
        path: Path (row from path-link array that indicates links used by path)
        initial_slot: Initial slot
        num_slots: Number of slots
        value: Value to set on link slot array

    Returns:
        Updated link slot array
    """
    return jax.vmap(set_path, in_axes=(0, 0, None, None, None))(link_slot_array, path, initial_slot, num_slots, value)

vmap_update_node_departure(node_departure_array, selected_nodes, value)

Called when implementing node action. Sets request departure time ("value") in place of first "inf" i.e. unoccupied index on node departure array for selected nodes.

Parameters:

Name	Type	Description	Default
node_departure_array	Array	(N x R) Node departure array	required
selected_nodes	Array	(N x 1) Selected nodes (non-zero value on selected node indices)	required
value	int	Value to set on node departure array	required

Returns:

Type	Description
Array	Updated node departure array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_update_node_departure(node_departure_array: chex.Array, selected_nodes: chex.Array, value: int) -> chex.Array:
    """Called when implementing node action.
    Sets request departure time ("value") in place of first "inf" i.e. unoccupied index on node departure array for selected nodes.

    Args:
        node_departure_array: (N x R) Node departure array
        selected_nodes: (N x 1) Selected nodes (non-zero value on selected node indices)
        value: Value to set on node departure array

    Returns:
        Updated node departure array
    """
    first_inf_indices = jnp.argmax(node_departure_array, axis=1).astype(MED_INT_DTYPE)
    return jax.vmap(update_selected_node_departure, in_axes=(0, 0, 0, None))(node_departure_array, selected_nodes, first_inf_indices, value)

vmap_update_node_resources(node_resource_array, selected_nodes)

Called when implementing node action. Sets requested node resources on selected nodes in place of first "zero" i.e. unoccupied index on node resource array for selected nodes.

Parameters:

Name	Type	Description	Default
node_resource_array		(N x R) Node resource array	required
selected_nodes		(N x 1) Requested resources on selected nodes	required

Returns:

Type	Description
	Updated node resource array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_update_node_resources(node_resource_array, selected_nodes):
    """Called when implementing node action.
    Sets requested node resources on selected nodes in place of first "zero" i.e.
    unoccupied index on node resource array for selected nodes.

    Args:
        node_resource_array: (N x R) Node resource array
        selected_nodes: (N x 1) Requested resources on selected nodes

    Returns:
        Updated node resource array
    """
    first_zero_indices = jnp.argmin(node_resource_array, axis=1)
    return jax.vmap(update_selected_node_resources, in_axes=(0, 0, 0))(node_resource_array, selected_nodes, first_zero_indices)

vmap_update_path_links(link_slot_array, path, initial_slot, num_slots, value)

Update relevant slots along links in path to current_val - val.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
path	Array	Path (row from path-link array that indicates links used by path)	required
initial_slot	int	Initial slot	required
num_slots	int	Number of slots	required
value	int	Value to subtract from link slot array	required

Returns:

Type	Description
Array	Updated link slot array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_update_path_links(link_slot_array: chex.Array, path: chex.Array, initial_slot: int, num_slots: int, value: int) -> chex.Array:
    """Update relevant slots along links in path to current_val - val.

    Args:
        link_slot_array: Link slot array
        path: Path (row from path-link array that indicates links used by path)
        initial_slot: Initial slot
        num_slots: Number of slots
        value: Value to subtract from link slot array

    Returns:
        Updated link slot array
    """
    return jax.vmap(update_path, in_axes=(0, 0, None, None, None))(link_slot_array, path, initial_slot, num_slots, value)

---

Models

ActorCriticGNN

Bases: Module

Combine the GNN actor and critic networks into a single class

Source code in xlron/models/models.py

class ActorCriticGNN(nn.Module):
    """Combine the GNN actor and critic networks into a single class"""
    activation: str = "tanh"
    num_layers: int = 2
    num_units: int = 64
    message_passing_steps: int = 1
    mlp_layers: int = None
    mlp_latent: int = None
    edge_embedding_size: int = 128
    edge_mlp_layers: int = 3
    edge_mlp_latent: int = 128
    edge_output_size_actor: int = 1
    edge_output_size_critic: int = 1
    global_embedding_size: int = 8
    global_mlp_layers: int = 0
    global_mlp_latent: int = 0
    global_output_size_actor: int = 0
    global_output_size_critic: int = 0
    node_embedding_size: int = 16
    node_mlp_layers: int = 2
    node_mlp_latent: int = 128
    node_output_size_actor: int = 0
    node_output_size_critic: int = 0
    attn_mlp_layers: int = 2
    attn_mlp_latent: int = 128
    gnn_mlp_layers: int = 1
    use_attention: bool = True
    normalise_by_link_length: bool = True  # Normalise the processed edge features by the link length
    gnn_layer_norm: bool = True
    mlp_layer_norm: bool = False
    vmap: bool = True
    temperature: float = 1.0
    # Launch power specific parameters
    min_power_dbm: float = 0.0
    max_power_dbm: float = 2.0
    step_power_dbm: float = 0.1
    discrete: bool = True
    # Beta distribution parameters (used only if discrete=False)
    min_concentration: float = 0.1
    max_concentration: float = 20.0
    epsilon: float = 1e-6
    # Bools to determine which actions to output
    output_path: bool = True
    output_power: bool = True
    assert edge_output_size_actor > 0
    assert edge_output_size_critic + global_output_size_critic > 0

    @property
    def num_power_levels(self):
        """Calculate number of power levels dynamically"""
        return int((self.max_power_dbm - self.min_power_dbm) / self.step_power_dbm) + 1

    @property
    def power_levels(self):
        """Calculate power levels dynamically"""
        return jnp.linspace(self.min_power_dbm, self.max_power_dbm, self.num_power_levels, dtype=LARGE_FLOAT_DTYPE)

    @nn.compact
    def __call__(self, state: EnvState, params: EnvParams):
        actor = ActorGNN(
            num_layers=self.num_layers,
            num_units=self.num_units,
            message_passing_steps=self.message_passing_steps,
            mlp_layers=self.mlp_layers,
            mlp_latent=self.mlp_latent,
            edge_embedding_size=self.edge_embedding_size,
            edge_mlp_layers=self.edge_mlp_layers,
            edge_mlp_latent=self.edge_mlp_latent,
            edge_output_size=self.edge_output_size_actor,
            global_embedding_size=self.global_embedding_size,
            global_mlp_layers=self.global_mlp_layers,
            global_mlp_latent=self.global_mlp_latent,
            global_output_size=self.global_output_size_actor,
            node_embedding_size=self.node_embedding_size,
            node_mlp_layers=self.node_mlp_layers,
            node_mlp_latent=self.node_mlp_latent,
            node_output_size=self.node_output_size_actor,
            attn_mlp_layers=self.attn_mlp_layers,
            attn_mlp_latent=self.attn_mlp_latent,
            use_attention=self.use_attention,
            normalise_by_link_length=self.normalise_by_link_length,
            gnn_layer_norm=self.gnn_layer_norm,
            mlp_layer_norm=self.mlp_layer_norm,
            temperature=self.temperature,
            min_power_dbm=self.min_power_dbm,
            max_power_dbm=self.max_power_dbm,
            step_power_dbm=self.step_power_dbm,
            discrete=self.discrete,
            min_concentration=self.min_concentration,
            max_concentration=self.max_concentration,
            epsilon=self.epsilon,
        )
        critic = CriticGNN(
            activation=self.activation,
            num_layers=self.num_layers,
            num_units=self.num_units,
            message_passing_steps=self.message_passing_steps,
            mlp_layers=self.mlp_layers,
            mlp_latent=self.mlp_latent,
            edge_embedding_size=self.edge_embedding_size,
            edge_mlp_layers=self.edge_mlp_layers,
            edge_mlp_latent=self.edge_mlp_latent,
            edge_output_size=self.edge_output_size_critic,
            global_embedding_size=self.global_embedding_size,
            global_mlp_layers=self.global_mlp_layers,
            global_mlp_latent=self.global_mlp_latent,
            global_output_size=self.global_output_size_critic,
            node_embedding_size=self.node_embedding_size,
            node_mlp_layers=self.node_mlp_layers,
            node_mlp_latent=self.node_mlp_latent,
            node_output_size=self.node_output_size_critic,
            attn_mlp_layers=self.attn_mlp_layers,
            attn_mlp_latent=self.attn_mlp_latent,
            use_attention=self.use_attention,
            normalise_by_link_length=self.normalise_by_link_length,
            gnn_layer_norm=self.gnn_layer_norm,
            mlp_layer_norm=self.mlp_layer_norm,
        )

        if self.vmap:
            actor = jax.vmap(actor, in_axes=(0, None))
            critic = jax.vmap(critic, in_axes=(0, None))
        actor_out = actor(state, params)
        critic_out = critic(state, params)
        return actor_out, critic_out

    def sample_action_path(self, seed, dist, log_prob=False, deterministic=False):
        """Sample an action from the distribution."""
        action = jnp.argmax(dist.probs()).astype(MED_INT_DTYPE) if deterministic else dist.sample(seed=seed)
        if log_prob:
            return action, dist.log_prob(action)
        return action

    def sample_action_power(self, seed, dist, log_prob=False, deterministic=False):
        """Sample an action and convert to power level"""
        if self.discrete:
            if deterministic:
                # Take most probable action
                raw_action = dist.mode()
            else:
                # Sample from distribution
                raw_action = dist.sample(seed=seed)
            processed_action = self.power_levels[raw_action]
        else:
            if deterministic:
                # Use mean of Beta distribution for deterministic action
                mean = dist.alpha / (dist.alpha + dist.beta)
                raw_action = jnp.clip(mean, self.epsilon, 1.0 - self.epsilon)
            else:
                # Sample from Beta (clipping to avoid edge values with undefined gradient) and scale to power range
                raw_action = jnp.clip(dist.sample(seed=seed), self.epsilon, 1.0 - self.epsilon)
            processed_action = self.min_power_dbm + raw_action * (self.max_power_dbm - self.min_power_dbm)
        processed_action = from_dbm(processed_action)
        if log_prob:
            return processed_action, dist.log_prob(raw_action)
        return processed_action

    def sample_action_path_power(self, seed, dist, log_prob=False, deterministic=False):
        """Sample an action from the distributions.
        This assumes dist is a tuple of path and power distributions."""
        path_action = self.sample_action_path(seed, dist[0], log_prob=log_prob, deterministic=deterministic)
        power_action  = self.sample_action_power(seed, dist[1], log_prob=log_prob, deterministic=deterministic)
        if log_prob:
            return path_action[0], power_action[0], path_action[1]+power_action[1]
        return path_action, power_action

    def sample_action(self, seed, dist, log_prob=False, deterministic=False):
        """Sample an action from the distributions.
        This assumes dist is a tuple of path and power distributions OR just the appropriate distribution."""
        if self.output_path and self.output_power:
            return self.sample_action_path_power(seed, dist, log_prob=log_prob, deterministic=deterministic)
        elif self.output_path:
            return self.sample_action_path(seed, dist, log_prob=log_prob, deterministic=deterministic)
        elif self.output_power:
            return self.sample_action_power(seed, dist, log_prob=log_prob, deterministic=deterministic)
        else:
            raise ValueError("No action type specified for sampling.")

num_power_levels property

Calculate number of power levels dynamically

power_levels property

Calculate power levels dynamically

sample_action(seed, dist, log_prob=False, deterministic=False)

Sample an action from the distributions. This assumes dist is a tuple of path and power distributions OR just the appropriate distribution.

Source code in xlron/models/models.py

def sample_action(self, seed, dist, log_prob=False, deterministic=False):
    """Sample an action from the distributions.
    This assumes dist is a tuple of path and power distributions OR just the appropriate distribution."""
    if self.output_path and self.output_power:
        return self.sample_action_path_power(seed, dist, log_prob=log_prob, deterministic=deterministic)
    elif self.output_path:
        return self.sample_action_path(seed, dist, log_prob=log_prob, deterministic=deterministic)
    elif self.output_power:
        return self.sample_action_power(seed, dist, log_prob=log_prob, deterministic=deterministic)
    else:
        raise ValueError("No action type specified for sampling.")

sample_action_path(seed, dist, log_prob=False, deterministic=False)

Sample an action from the distribution.

Source code in xlron/models/models.py

def sample_action_path(self, seed, dist, log_prob=False, deterministic=False):
    """Sample an action from the distribution."""
    action = jnp.argmax(dist.probs()).astype(MED_INT_DTYPE) if deterministic else dist.sample(seed=seed)
    if log_prob:
        return action, dist.log_prob(action)
    return action

sample_action_path_power(seed, dist, log_prob=False, deterministic=False)

Sample an action from the distributions. This assumes dist is a tuple of path and power distributions.

Source code in xlron/models/models.py

def sample_action_path_power(self, seed, dist, log_prob=False, deterministic=False):
    """Sample an action from the distributions.
    This assumes dist is a tuple of path and power distributions."""
    path_action = self.sample_action_path(seed, dist[0], log_prob=log_prob, deterministic=deterministic)
    power_action  = self.sample_action_power(seed, dist[1], log_prob=log_prob, deterministic=deterministic)
    if log_prob:
        return path_action[0], power_action[0], path_action[1]+power_action[1]
    return path_action, power_action

sample_action_power(seed, dist, log_prob=False, deterministic=False)

Sample an action and convert to power level

Source code in xlron/models/models.py

def sample_action_power(self, seed, dist, log_prob=False, deterministic=False):
    """Sample an action and convert to power level"""
    if self.discrete:
        if deterministic:
            # Take most probable action
            raw_action = dist.mode()
        else:
            # Sample from distribution
            raw_action = dist.sample(seed=seed)
        processed_action = self.power_levels[raw_action]
    else:
        if deterministic:
            # Use mean of Beta distribution for deterministic action
            mean = dist.alpha / (dist.alpha + dist.beta)
            raw_action = jnp.clip(mean, self.epsilon, 1.0 - self.epsilon)
        else:
            # Sample from Beta (clipping to avoid edge values with undefined gradient) and scale to power range
            raw_action = jnp.clip(dist.sample(seed=seed), self.epsilon, 1.0 - self.epsilon)
        processed_action = self.min_power_dbm + raw_action * (self.max_power_dbm - self.min_power_dbm)
    processed_action = from_dbm(processed_action)
    if log_prob:
        return processed_action, dist.log_prob(raw_action)
    return processed_action

ActorCriticMLP

Bases: Module

Source code in xlron/models/models.py

class ActorCriticMLP(nn.Module):
    action_dim: Sequence[int]
    activation: str = "tanh"
    num_layers: int = 2
    num_units: int = 64
    layer_norm: bool = False
    temperature: float = 1.0

    @nn.compact
    def __call__(self, x):
        actor_mean = make_dense_layers(x, self.num_units, self.num_layers, self.activation)
        stacked_logits = []
        actor_mean_dim = nn.Dense(
            self.action_dim,
            kernel_init=orthogonal(0.01),
            bias_init=constant(0.0),
            dtype=COMPUTE_DTYPE,
            param_dtype=PARAMS_DTYPE,
        )(actor_mean)
        logits = actor_mean_dim / self.temperature
        stacked_logits.append(logits)

        # If there are multiple action dimensions, concatenate the logits
        action_dist = distrax.Categorical(logits=logits)

        critic = make_dense_layers(x, self.num_units, self.num_layers, self.activation, layer_norm=self.layer_norm)
        critic = nn.Dense(
            1,
            kernel_init=orthogonal(1.0),
            bias_init=constant(0.0),
            dtype=COMPUTE_DTYPE,
            param_dtype=PARAMS_DTYPE,
        )(
            critic
        )

        return action_dist, jnp.squeeze(critic, axis=-1)

    def sample_action(self, seed,  dist, log_prob=False, deterministic=False):
        """Sample an action from the distribution"""
        action = jnp.argmax(dist.probs()) if deterministic else dist.sample(seed=seed)
        if log_prob:
            return action, dist.log_prob(action)
        return action

sample_action(seed, dist, log_prob=False, deterministic=False)

Sample an action from the distribution

Source code in xlron/models/models.py

def sample_action(self, seed,  dist, log_prob=False, deterministic=False):
    """Sample an action from the distribution"""
    action = jnp.argmax(dist.probs()) if deterministic else dist.sample(seed=seed)
    if log_prob:
        return action, dist.log_prob(action)
    return action

ActorGNN

Bases: Module

Actor network for PPO algorithm. Takes the current state and returns a distrax.Categorical distribution over actions.

Source code in xlron/models/models.py

class ActorGNN(nn.Module):
    """
    Actor network for PPO algorithm. Takes the current state and returns a distrax.Categorical distribution
    over actions.
    """
    activation: str = "tanh"
    num_layers: int = 2
    num_units: int = 64
    mlp_layers: int = None
    mlp_latent: int = None
    edge_embedding_size: int = 128
    edge_mlp_layers: int = 3
    edge_mlp_latent: int = 128
    edge_output_size: int = 0
    global_embedding_size: int = 8
    global_mlp_layers: int = 0
    global_mlp_latent: int = 0
    global_output_size: int = 0
    node_embedding_size: int = 16
    node_mlp_layers: int = 2
    node_mlp_latent: int = 128
    node_output_size: int = 0
    attn_mlp_layers: int = 2
    attn_mlp_latent: int = 128
    dropout_rate: float = 0
    deterministic: bool = False
    message_passing_steps: int = 1
    use_attention: bool = True
    normalise_by_link_length: bool = True  # Normalise the processed edge features by the link length
    gnn_layer_norm: bool = True
    mlp_layer_norm: bool = False
    temperature: float = 1.0
    # Launch power specific parameters
    min_power_dbm: float = 0.0
    max_power_dbm: float = 2.0
    step_power_dbm: float = 0.1
    discrete: bool = True
    # Beta distribution parameters (used only if discrete=False)
    min_concentration: float = 0.1
    max_concentration: float = 20.0
    epsilon: float = 1e-6

    @property
    def num_power_levels(self):
        """Calculate number of power levels dynamically"""
        return int((self.max_power_dbm - self.min_power_dbm) / self.step_power_dbm) + 1

    @property
    def power_levels(self):
        """Calculate power levels dynamically"""
        return jnp.linspace(self.min_power_dbm, self.max_power_dbm, self.num_power_levels, dtype=SMALL_FLOAT_DTYPE)

    @nn.compact
    def __call__(self, state: EnvState, params: EnvParams):
        """
        The ActorGNN network takes the current network state in the form of a GraphTuple and returns
        a distrax.Categorical distribution over actions.
        The graph is processed by a GraphNet module, and the resulting graph is indexed to construct a matrix of
        the edge features. The edge features are then normalised by the link length from the environment parameters,
        and the current request is read from the request array.
        The request is used to retrieve the edge features from the edge_features array for the corresponding
        shortest k-paths. The edge features are aggregated for each path according to the "agg_func" e.g. sum,
        and the action distribution array is updated.
        Returns a distrax.Categorical distribution, from which actions can be sampled.

        :param state: EnvState
        :param params: EnvParams

        :return: distrax.Categorical distribution over actions
        """
        processed_graph = GraphNet(
            message_passing_steps=self.message_passing_steps,
            mlp_layers=self.mlp_layers,
            mlp_latent=self.mlp_latent,
            edge_embedding_size=self.edge_embedding_size,
            edge_mlp_layers=self.edge_mlp_layers,
            edge_mlp_latent=self.edge_mlp_latent,
            edge_output_size=self.edge_output_size,
            global_embedding_size=self.global_embedding_size,
            global_mlp_layers=self.global_mlp_layers,
            global_mlp_latent=self.global_mlp_latent,
            global_output_size=self.global_output_size,
            node_embedding_size=self.node_embedding_size,
            node_mlp_layers=self.node_mlp_layers,
            node_mlp_latent=self.node_mlp_latent,
            node_output_size=self.node_output_size,
            attn_mlp_layers=self.attn_mlp_layers,
            attn_mlp_latent=self.attn_mlp_latent,
            use_attention=self.use_attention,
            gnn_layer_norm=self.gnn_layer_norm,
            mlp_layer_norm=self.mlp_layer_norm,
        )(state.graph)

        # Index edge features to resemble the link-slot array
        edge_features = processed_graph.edges if params.directed_graph else processed_graph.edges[:len(processed_graph.edges) // 2]
        # Normalise features by normalised link length from state.link_length_array
        if self.normalise_by_link_length:
            edge_features = edge_features * (params.link_length_array.val/jnp.sum(params.link_length_array.val, promote_integers=False))
        # Get the current request and initialise array of action distributions per path
        nodes_sd, requested_bw = read_rsa_request(state.request_array)
        init_action_array = jnp.zeros(params.k_paths * self.edge_output_size, dtype=SMALL_INT_DTYPE)

        # Define a body func to retrieve path slots and update action array
        def get_path_action_dist(i, action_array):
            # Get the processed graph edge features corresponding to the i-th path
            path_features = get_path_slots(edge_features, params, nodes_sd, i, agg_func="sum")
            # Update the action array with the path features
            action_array = jax.lax.dynamic_update_slice(action_array, path_features, (i * self.edge_output_size,))
            return action_array

        path_action_logits = jax.lax.fori_loop(0, params.k_paths, get_path_action_dist, init_action_array)
        path_action_logits = jnp.reshape(path_action_logits, (-1,)) / self.temperature

        power_action_dist = None
        mlp_features = [self.num_units] * self.num_layers
        output_size = self.num_power_levels if self.discrete else 2
        path_mlp = MLP(
            mlp_features + [output_size],
            dropout_rate=self.dropout_rate,
            deterministic=self.deterministic,
            layer_norm=self.mlp_layer_norm,
        )
        if params.__class__.__name__ == "RSAGNModelEnvParams":
            if self.global_output_size > 0:
                power_logits = processed_graph.globals.reshape((-1,)) / self.temperature
            else:
                init_feature_array = jnp.zeros((params.k_paths, edge_features.shape[1]), dtype=LARGE_FLOAT_DTYPE)
                # Define a body func to retrieve path slots and update action array
                def get_power_action_dist(i, feature_array):
                    # Get the processed graph edge features corresponding to the i-th path
                    path_features = get_path_slots(edge_features, params, nodes_sd, i, agg_func="sum").reshape((1, -1))
                    # Update the array with the path features
                    action_array = jax.lax.dynamic_update_slice(feature_array, path_features,(i, 0))
                    return action_array
                path_feature_batch = jax.lax.fori_loop(0, params.k_paths, get_power_action_dist, init_feature_array)
                power_logits = path_mlp(path_feature_batch)
            if self.discrete:
                power_action_dist = distrax.Categorical(logits=power_logits)
            else:
                alpha = self.min_concentration + jax.nn.softplus(power_logits) * (
                        self.max_concentration - self.min_concentration
                )
                beta = self.min_concentration + jax.nn.softplus(power_logits) * (
                        self.max_concentration - self.min_concentration
                )
                power_action_dist = distrax.Beta(alpha, beta)

        # Return a distrax.Categorical distribution over actions (which can be masked later)
        path_action_dist = distrax.Categorical(logits=path_action_logits)
        return (path_action_dist, power_action_dist)

num_power_levels property

Calculate number of power levels dynamically

power_levels property

Calculate power levels dynamically

__call__(state, params)

The ActorGNN network takes the current network state in the form of a GraphTuple and returns a distrax.Categorical distribution over actions. The graph is processed by a GraphNet module, and the resulting graph is indexed to construct a matrix of the edge features. The edge features are then normalised by the link length from the environment parameters, and the current request is read from the request array. The request is used to retrieve the edge features from the edge_features array for the corresponding shortest k-paths. The edge features are aggregated for each path according to the "agg_func" e.g. sum, and the action distribution array is updated. Returns a distrax.Categorical distribution, from which actions can be sampled.

:param state: EnvState :param params: EnvParams

:return: distrax.Categorical distribution over actions

Source code in xlron/models/models.py

@nn.compact
def __call__(self, state: EnvState, params: EnvParams):
    """
    The ActorGNN network takes the current network state in the form of a GraphTuple and returns
    a distrax.Categorical distribution over actions.
    The graph is processed by a GraphNet module, and the resulting graph is indexed to construct a matrix of
    the edge features. The edge features are then normalised by the link length from the environment parameters,
    and the current request is read from the request array.
    The request is used to retrieve the edge features from the edge_features array for the corresponding
    shortest k-paths. The edge features are aggregated for each path according to the "agg_func" e.g. sum,
    and the action distribution array is updated.
    Returns a distrax.Categorical distribution, from which actions can be sampled.

    :param state: EnvState
    :param params: EnvParams

    :return: distrax.Categorical distribution over actions
    """
    processed_graph = GraphNet(
        message_passing_steps=self.message_passing_steps,
        mlp_layers=self.mlp_layers,
        mlp_latent=self.mlp_latent,
        edge_embedding_size=self.edge_embedding_size,
        edge_mlp_layers=self.edge_mlp_layers,
        edge_mlp_latent=self.edge_mlp_latent,
        edge_output_size=self.edge_output_size,
        global_embedding_size=self.global_embedding_size,
        global_mlp_layers=self.global_mlp_layers,
        global_mlp_latent=self.global_mlp_latent,
        global_output_size=self.global_output_size,
        node_embedding_size=self.node_embedding_size,
        node_mlp_layers=self.node_mlp_layers,
        node_mlp_latent=self.node_mlp_latent,
        node_output_size=self.node_output_size,
        attn_mlp_layers=self.attn_mlp_layers,
        attn_mlp_latent=self.attn_mlp_latent,
        use_attention=self.use_attention,
        gnn_layer_norm=self.gnn_layer_norm,
        mlp_layer_norm=self.mlp_layer_norm,
    )(state.graph)

    # Index edge features to resemble the link-slot array
    edge_features = processed_graph.edges if params.directed_graph else processed_graph.edges[:len(processed_graph.edges) // 2]
    # Normalise features by normalised link length from state.link_length_array
    if self.normalise_by_link_length:
        edge_features = edge_features * (params.link_length_array.val/jnp.sum(params.link_length_array.val, promote_integers=False))
    # Get the current request and initialise array of action distributions per path
    nodes_sd, requested_bw = read_rsa_request(state.request_array)
    init_action_array = jnp.zeros(params.k_paths * self.edge_output_size, dtype=SMALL_INT_DTYPE)

    # Define a body func to retrieve path slots and update action array
    def get_path_action_dist(i, action_array):
        # Get the processed graph edge features corresponding to the i-th path
        path_features = get_path_slots(edge_features, params, nodes_sd, i, agg_func="sum")
        # Update the action array with the path features
        action_array = jax.lax.dynamic_update_slice(action_array, path_features, (i * self.edge_output_size,))
        return action_array

    path_action_logits = jax.lax.fori_loop(0, params.k_paths, get_path_action_dist, init_action_array)
    path_action_logits = jnp.reshape(path_action_logits, (-1,)) / self.temperature

    power_action_dist = None
    mlp_features = [self.num_units] * self.num_layers
    output_size = self.num_power_levels if self.discrete else 2
    path_mlp = MLP(
        mlp_features + [output_size],
        dropout_rate=self.dropout_rate,
        deterministic=self.deterministic,
        layer_norm=self.mlp_layer_norm,
    )
    if params.__class__.__name__ == "RSAGNModelEnvParams":
        if self.global_output_size > 0:
            power_logits = processed_graph.globals.reshape((-1,)) / self.temperature
        else:
            init_feature_array = jnp.zeros((params.k_paths, edge_features.shape[1]), dtype=LARGE_FLOAT_DTYPE)
            # Define a body func to retrieve path slots and update action array
            def get_power_action_dist(i, feature_array):
                # Get the processed graph edge features corresponding to the i-th path
                path_features = get_path_slots(edge_features, params, nodes_sd, i, agg_func="sum").reshape((1, -1))
                # Update the array with the path features
                action_array = jax.lax.dynamic_update_slice(feature_array, path_features,(i, 0))
                return action_array
            path_feature_batch = jax.lax.fori_loop(0, params.k_paths, get_power_action_dist, init_feature_array)
            power_logits = path_mlp(path_feature_batch)
        if self.discrete:
            power_action_dist = distrax.Categorical(logits=power_logits)
        else:
            alpha = self.min_concentration + jax.nn.softplus(power_logits) * (
                    self.max_concentration - self.min_concentration
            )
            beta = self.min_concentration + jax.nn.softplus(power_logits) * (
                    self.max_concentration - self.min_concentration
            )
            power_action_dist = distrax.Beta(alpha, beta)

    # Return a distrax.Categorical distribution over actions (which can be masked later)
    path_action_dist = distrax.Categorical(logits=path_action_logits)
    return (path_action_dist, power_action_dist)

GraphNet

Bases: Module

A complete Graph Network model defined with Jraph.

Source code in xlron/models/models.py

class GraphNet(nn.Module):
    """A complete Graph Network model defined with Jraph."""

    message_passing_steps: int
    mlp_layers: int = None
    mlp_latent: int = None
    edge_embedding_size: int = 128
    edge_mlp_layers: int = 3
    edge_mlp_latent: int = 128
    edge_output_size: int = 0
    global_embedding_size: int = 8
    global_mlp_layers: int = 0
    global_mlp_latent: int = 0
    global_output_size: int = 0
    node_embedding_size: int = 16
    node_mlp_layers: int = 2
    node_mlp_latent: int = 128
    node_output_size: int = 0
    attn_mlp_layers: int = 2
    attn_mlp_latent: int = 128
    dropout_rate: float = 0
    skip_connections: bool = True
    use_edge_model: bool = True
    gnn_layer_norm: bool = True
    mlp_layer_norm: bool = False
    deterministic: bool = True  # If true, no dropout (better for RL purposes)
    use_attention: bool = True

    @nn.compact
    def __call__(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:
        # Template code from here: https://github.com/google/flax/blob/main/examples/ogbg_molpcba/models.py
        # We will first linearly project the original features as 'embeddings'.
        if self.mlp_latent is not None:
            global_mlp_dims = edge_mlp_dims = node_mlp_dims = attn_mlp_dims = [self.mlp_latent] * self.mlp_layers
        else:
            global_mlp_dims = [self.global_mlp_latent] * self.global_mlp_layers
            edge_mlp_dims = [self.edge_mlp_latent] * self.edge_mlp_layers
            node_mlp_dims = [self.node_mlp_latent] * self.node_mlp_layers
            attn_mlp_dims = [self.attn_mlp_latent] * self.attn_mlp_layers
        if self.skip_connections:
            # If using skip connections, we need to add the input dimensions to the output dimensions,
            # so that the output of the MLP is the same size as the input for summing output/input
            edge_mlp_dims = edge_mlp_dims + [self.edge_embedding_size]
            node_mlp_dims = node_mlp_dims + [self.node_embedding_size]
            global_mlp_dims = global_mlp_dims + [self.global_embedding_size]

        embedder = jraph.GraphMapFeatures(
            embed_edge_fn=nn.Dense(self.edge_embedding_size),
            embed_node_fn=nn.Dense(self.node_embedding_size),
            embed_global_fn=nn.Dense(self.global_embedding_size),
        )
        if graphs.edges.ndim >= 3:
            # Dims are (edges, slots, features e.g. power, source/dest)
            # Keep the leading dimension fixed and combine the remaining dimensions
            edges = graphs.edges.reshape((graphs.edges.shape[0], -1))
            graphs = graphs._replace(edges=edges)
        processed_graphs = embedder(graphs)

        # Now, we will apply a Graph Network once for each message-passing round.
        for _ in range(self.message_passing_steps):
            if self.use_edge_model:
                update_edge_fn = jraph.concatenated_args(
                    MLP(
                        edge_mlp_dims,
                        dropout_rate=self.dropout_rate,
                        deterministic=self.deterministic,
                        layer_norm=self.mlp_layer_norm,
                    )
                )
            else:
                update_edge_fn = None

            update_node_fn = jraph.concatenated_args(
                MLP(
                    node_mlp_dims,
                    dropout_rate=self.dropout_rate,
                    deterministic=self.deterministic,
                    layer_norm=self.mlp_layer_norm,
                )
            )
            if self.global_output_size > 0:
                update_global_fn = jraph.concatenated_args(
                    MLP(
                        global_mlp_dims,
                        dropout_rate=self.dropout_rate,
                        deterministic=self.deterministic,
                        layer_norm=self.mlp_layer_norm,
                    )
                )
            else:
                update_global_fn = None

            if self.use_attention:
                def _attention_logit_fn(edges, sender_attr, receiver_attr, global_edge_attributes):
                    """Calculate attention logits using edges, nodes and global attributes."""
                    x = jnp.concatenate((edges, sender_attr, receiver_attr, global_edge_attributes), axis=1)
                    return MLP(attn_mlp_dims + [1], dropout_rate=self.dropout_rate,
                               deterministic=self.deterministic)(x)

                def _attention_reduce_fn(
                    edges: jnp.ndarray, attention: jnp.ndarray
                ) -> jnp.ndarray:
                    # TODO - try more sophisticated attention reduce function (not sure what it would be)
                    #  (here might help https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial7/GNN_overview.html)
                    return attention * edges

                graph_net = GraphNetGAT(
                    update_node_fn=update_node_fn,
                    update_edge_fn=update_edge_fn,  # Update the edges with MLP prior to attention
                    update_global_fn=update_global_fn,
                    attention_logit_fn=_attention_logit_fn,
                    attention_reduce_fn=_attention_reduce_fn,
                )
            else:
                graph_net = GraphNetwork(
                    update_node_fn=update_node_fn,
                    update_edge_fn=update_edge_fn,
                    update_global_fn=update_global_fn,
                )

            if self.skip_connections:
                processed_graphs = add_graphs_tuples(
                graph_net(processed_graphs), processed_graphs
            )
            else:
                processed_graphs = graph_net(processed_graphs)

            if self.gnn_layer_norm:
                processed_graphs = processed_graphs._replace(
                    nodes=nn.LayerNorm(dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)(processed_graphs.nodes),
                    edges=nn.LayerNorm(dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)(processed_graphs.edges),
                    globals=nn.LayerNorm(dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)(processed_graphs.globals) if processed_graphs.globals is not None else None,
                )

        decoder = jraph.GraphMapFeatures(
            embed_global_fn=nn.Dense(self.global_output_size, dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,) if self.global_output_size > 0 else None,
            embed_node_fn=nn.Dense(self.node_output_size, dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,) if self.node_output_size > 0 else None,
            embed_edge_fn=nn.Dense(self.edge_output_size, dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,),
        )
        processed_graphs = decoder(processed_graphs)

        return processed_graphs

LaunchPowerActorCriticMLP

Bases: Module

In this implementation, we take an observation of th current request + statistics on each of the K candidate paths. We make K forward passes, one for each path, and output a distribution over power levels for each path. In action selection, we then sample from each distribution and use the sampled power levels to mask paths for the routing heuristic, which then determines the path taken. The selected path index is then used to select which output action, distribution, and value to use for the loss calculation.

Source code in xlron/models/models.py

class LaunchPowerActorCriticMLP(nn.Module):
    """In this implementation, we take an observation of th current request + statistics on each of the K candidate paths.
    We make K forward passes, one for each path, and output a distribution over power levels for each path.
    In action selection, we then sample from each distribution and use the sampled power levels to mask paths for the
    routing heuristic, which then determines the path taken. The selected path index is then used to select which output
    action, distribution, and value to use for the loss calculation."""
    action_dim: Sequence[int]
    activation: str = "tanh"
    num_layers: int = 2
    num_units: int = 64
    layer_norm: bool = False
    min_power_dbm: float = 0.0
    max_power_dbm: float = 2.0
    step_power_dbm: float = 0.1
    discrete: bool = True
    temperature: float = 1.0
    k_paths: int = 5
    num_base_features: int = 4
    num_path_features: int = 7
    # Beta distribution parameters (used only if discrete=False)
    min_concentration: float = 0.1
    max_concentration: float = 20.0
    epsilon: float = 1e-6

    @property
    def num_power_levels(self):
        """Calculate number of power levels dynamically"""
        return int((self.max_power_dbm - self.min_power_dbm) / self.step_power_dbm) + 1

    @property
    def power_levels(self):
        """Calculate power levels dynamically"""
        return jnp.linspace(self.min_power_dbm, self.max_power_dbm, self.num_power_levels, dtype=SMALL_FLOAT_DTYPE)

    @nn.compact
    def __call__(self, x):
        # Helper function to create MLP layers
        def make_mlp(prefix):
            layers = []
            for i in range(self.num_layers):
                layers.append(nn.Dense(
                    self.num_units,
                    kernel_init=orthogonal(np.sqrt(2)),
                    name=f"{prefix}_dense_{i}",
                    dtype=COMPUTE_DTYPE,
                    param_dtype=PARAMS_DTYPE,
                ))
                if self.layer_norm:
                    layers.append(nn.LayerNorm(name=f"{prefix}_norm_{i}", dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,))
            return layers

        # Initialize actor network layers
        actor_net = make_mlp("actor")
        if self.discrete:
            actor_out = nn.Dense(
                self.num_power_levels,
                kernel_init=orthogonal(0.01),
                name="actor_output",
                dtype=COMPUTE_DTYPE,
                param_dtype=PARAMS_DTYPE,
            )
        else:
            alpha_out = nn.Dense(1, kernel_init=orthogonal(0.01), name="alpha", dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)
            beta_out = nn.Dense(1, kernel_init=orthogonal(0.01), name="beta", dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)

        def activate(x):
            if self.activation == "relu": return jax.nn.relu(x)
            if self.activation == "crelu": return crelu(x)
            return jnp.tanh(x)

        def forward(x, layers):
            for layer in layers:
                x = layer(x)
                if isinstance(layer, nn.Dense):
                    x = activate(x)
            return x

        num_base_features = self.num_base_features
        num_path_features = self.num_path_features
        temperature = self.temperature
        discrete = self.discrete
        min_concentration = self.min_concentration
        max_concentration = self.max_concentration

        # Create a class to handle the scan
        class PathProcessor(nn.Module):
            def __call__(self, carry, i):
                base = x[:num_base_features]
                path = jax.lax.dynamic_slice(
                    x,
                    (num_base_features + i *num_path_features,),
                    (num_path_features,)
                )
                features = jnp.concatenate([base, path])
                actor_hidden = forward(features, actor_net)
                if discrete:
                    out = actor_out(actor_hidden) / temperature
                else:
                    alpha = min_concentration + jax.nn.softplus(alpha_out(actor_hidden)) * (
                            max_concentration - min_concentration
                    )
                    beta = min_concentration + jax.nn.softplus(beta_out(actor_hidden)) * (
                            max_concentration - min_concentration
                    )
                    out = (alpha, beta)

                return carry, out

        # Scan over paths
        _, dist_params = nn.scan(
            PathProcessor,
            variable_broadcast="params",
            split_rngs={"params": False},
            length=self.k_paths,
        )()(None, jnp.arange(self.k_paths, dtype=MED_INT_DTYPE))

        # Initialize critic network layers
        critic_net = make_mlp("critic")
        critic_out = nn.Dense(1, kernel_init=orthogonal(1.0), name="critic_output", dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)
        value = jnp.squeeze(critic_out(forward(x, critic_net)), axis=-1)

        # Create appropriate distribution
        if self.discrete:
            dist = distrax.Categorical(logits=dist_params)
        else:
            dist = distrax.Beta(dist_params[0].reshape((self.k_paths,)), dist_params[1].reshape((self.k_paths,)))
        # N.B. that this is a single distribution object, but it is batched over K paths

        return [None, dist], value

    def sample_action(self, seed, dist, log_prob=False, deterministic=False):
        """Sample an action and convert to power level"""
        if self.discrete:
            if deterministic:
                # Take most probable action
                raw_action = dist.mode()
            else:
                # Sample from distribution
                raw_action = dist.sample(seed=seed)
            processed_action = self.power_levels[raw_action].reshape((self.k_paths, 1))
        else:
            if deterministic:
                # Use mean of Beta distribution for deterministic action
                mean = dist.alpha / (dist.alpha + dist.beta)
                raw_action = jnp.clip(mean, self.epsilon, 1.0 - self.epsilon)
            else:
                # Sample from Beta (clipping to avoid edge values with undefined gradient) and scale to power range
                raw_action = jnp.clip(dist.sample(seed=seed), self.epsilon, 1.0 - self.epsilon)
            processed_action = self.min_power_dbm + raw_action * (self.max_power_dbm - self.min_power_dbm)
        processed_action = from_dbm(processed_action)
        if log_prob:
            return processed_action, dist.log_prob(jnp.squeeze(raw_action))
        return processed_action

    def get_action_probs(self, dist):
        """Get probabilities for discrete case or pdf for continuous case"""
        if self.discrete:
            return dist.probs()
        else:
            x = jnp.linspace(0, 1, 100)
            return dist.prob(x)

num_power_levels property

Calculate number of power levels dynamically

power_levels property

Calculate power levels dynamically

get_action_probs(dist)

Get probabilities for discrete case or pdf for continuous case

Source code in xlron/models/models.py

def get_action_probs(self, dist):
    """Get probabilities for discrete case or pdf for continuous case"""
    if self.discrete:
        return dist.probs()
    else:
        x = jnp.linspace(0, 1, 100)
        return dist.prob(x)

sample_action(seed, dist, log_prob=False, deterministic=False)

Sample an action and convert to power level

Source code in xlron/models/models.py

def sample_action(self, seed, dist, log_prob=False, deterministic=False):
    """Sample an action and convert to power level"""
    if self.discrete:
        if deterministic:
            # Take most probable action
            raw_action = dist.mode()
        else:
            # Sample from distribution
            raw_action = dist.sample(seed=seed)
        processed_action = self.power_levels[raw_action].reshape((self.k_paths, 1))
    else:
        if deterministic:
            # Use mean of Beta distribution for deterministic action
            mean = dist.alpha / (dist.alpha + dist.beta)
            raw_action = jnp.clip(mean, self.epsilon, 1.0 - self.epsilon)
        else:
            # Sample from Beta (clipping to avoid edge values with undefined gradient) and scale to power range
            raw_action = jnp.clip(dist.sample(seed=seed), self.epsilon, 1.0 - self.epsilon)
        processed_action = self.min_power_dbm + raw_action * (self.max_power_dbm - self.min_power_dbm)
    processed_action = from_dbm(processed_action)
    if log_prob:
        return processed_action, dist.log_prob(jnp.squeeze(raw_action))
    return processed_action

MLP

Bases: Module

A multi-layer perceptron.

Source code in xlron/models/models.py

class MLP(nn.Module):
    """A multi-layer perceptron."""

    feature_sizes: Sequence[int]
    dropout_rate: float = 0
    deterministic: bool = True
    activation: Callable[[jnp.ndarray], jnp.ndarray] = nn.tanh
    layer_norm: bool = False

    @nn.compact
    def __call__(self, inputs):
        x = inputs
        for size in self.feature_sizes:
            x = nn.Dense(features=size,
                         dtype=COMPUTE_DTYPE,
                         param_dtype=PARAMS_DTYPE,)(x)
            x = self.activation(x)
            x = nn.Dropout(rate=self.dropout_rate, deterministic=self.deterministic)(
                x
            )
            if self.layer_norm:
                x = nn.LayerNorm(dtype=COMPUTE_DTYPE, param_dtype=PARAMS_DTYPE,)(x)
        return x

add_graphs_tuples(graphs, other_graphs)

Adds the nodes, edges and global features from other_graphs to graphs.

Source code in xlron/models/models.py

def add_graphs_tuples(
    graphs: jraph.GraphsTuple, other_graphs: jraph.GraphsTuple
) -> jraph.GraphsTuple:
    """Adds the nodes, edges and global features from other_graphs to graphs."""
    return graphs._replace(
        nodes=graphs.nodes + other_graphs.nodes,
        edges=graphs.edges + other_graphs.edges,
        globals=graphs.globals + other_graphs.globals if graphs.globals is not None else None,
    )

crelu(x)

Computes the Concatenated ReLU (CReLU) activation function.

Source code in xlron/models/models.py

def crelu(x):
    """Computes the Concatenated ReLU (CReLU) activation function."""
    x = jnp.concatenate([x, -x], axis=-1)
    return nn.relu(x)

---

Heuristics

DeepRMSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for DeepRMSA.

Parameters:

Name	Type	Description	Default
path_stats	Array	Path stats array containing	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class DeepRMSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for DeepRMSA.

    Args:
        path_stats (chex.Array): Path stats array containing
        1. Required slots on path
        2. Total available slots on path
        3. Size of 1st free spectrum block
        4. Avg. free block size
    """
    path_stats: chex.Array

EnvParams

Dataclass to hold environment parameters. Parameters are immutable.

Parameters:

Name	Type	Description	Default
max_requests	Scalar	Maximum number of requests in an episode	required
incremental_loading	Scalar	Incremental increase in traffic load (non-expiring requests)	required
end_first_blocking	Scalar	End episode on first blocking event	required
continuous_operation	Scalar	If True, do not reset the environment at the end of an episode	required
edges	Array	Two column array defining source-dest node-pair edges of the graph	required
slot_size	Scalar	Spectral width of frequency slot in GHz	required
consider_modulation_format	Scalar	If True, consider modulation format to determine required slots	required
link_length_array	Array	Array of link lengths	required
aggregate_slots	Scalar	Number of slots to aggregate into a single action (First-Fit with aggregation)	required
guardband	Scalar	Guard band in slots	required
directed_graph	bool	Whether graph is directed (one fibre per link per transmission direction)	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvParams:
    """Dataclass to hold environment parameters. Parameters are immutable.

    Args:
        max_requests (chex.Scalar): Maximum number of requests in an episode
        incremental_loading (chex.Scalar): Incremental increase in traffic load (non-expiring requests)
        end_first_blocking (chex.Scalar): End episode on first blocking event
        continuous_operation (chex.Scalar): If True, do not reset the environment at the end of an episode
        edges (chex.Array): Two column array defining source-dest node-pair edges of the graph
        slot_size (chex.Scalar): Spectral width of frequency slot in GHz
        consider_modulation_format (chex.Scalar): If True, consider modulation format to determine required slots
        link_length_array (chex.Array): Array of link lengths
        aggregate_slots (chex.Scalar): Number of slots to aggregate into a single action (First-Fit with aggregation)
        guardband (chex.Scalar): Guard band in slots
        directed_graph (bool): Whether graph is directed (one fibre per link per transmission direction)
    """
    max_requests: chex.Scalar = struct.field(pytree_node=False)
    incremental_loading: chex.Scalar = struct.field(pytree_node=False)
    end_first_blocking: chex.Scalar = struct.field(pytree_node=False)
    continuous_operation: chex.Scalar = struct.field(pytree_node=False)
    edges: chex.Array = struct.field(pytree_node=False)
    slot_size: chex.Scalar = struct.field(pytree_node=False)
    consider_modulation_format: chex.Scalar = struct.field(pytree_node=False)
    link_length_array: chex.Array = struct.field(pytree_node=False)
    aggregate_slots: chex.Scalar = struct.field(pytree_node=False)
    guardband: chex.Scalar = struct.field(pytree_node=False)
    directed_graph: bool = struct.field(pytree_node=False)
    maximise_throughput: bool = struct.field(pytree_node=False)
    reward_type: str = struct.field(pytree_node=False)
    values_bw: chex.Array = struct.field(pytree_node=False)
    truncate_holding_time: bool = struct.field(pytree_node=False)
    traffic_array: bool = struct.field(pytree_node=False)
    pack_path_bits: bool = struct.field(pytree_node=False)
    relative_arrival_times: bool = struct.field(pytree_node=False)

EnvState

Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

Parameters:

Name	Type	Description	Default
current_time	Scalar	Current time in environment	required
holding_time	Scalar	Holding time of current request	required
total_timesteps	Scalar	Total timesteps in environment	required
total_requests	Scalar	Total requests in environment	required
graph	GraphsTuple	Graph tuple representing network state	required
full_link_slot_mask	Array	Action mask for link slot action (including if slot actions are aggregated)	required
accepted_services	Array	Number of accepted services	required
accepted_bitrate	Array	Accepted bitrate	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class EnvState:
    """Dataclass to hold environment state. State is mutable and arrays are traced on JIT compilation.

    Args:
        current_time (chex.Scalar): Current time in environment
        holding_time (chex.Scalar): Holding time of current request
        total_timesteps (chex.Scalar): Total timesteps in environment
        total_requests (chex.Scalar): Total requests in environment
        graph (jraph.GraphsTuple): Graph tuple representing network state
        full_link_slot_mask (chex.Array): Action mask for link slot action (including if slot actions are aggregated)
        accepted_services (chex.Array): Number of accepted services
        accepted_bitrate (chex.Array): Accepted bitrate
        """
    current_time: chex.Scalar
    holding_time: chex.Scalar
    arrival_time: chex.Scalar
    total_timesteps: chex.Scalar
    total_requests: chex.Scalar
    graph: jraph.GraphsTuple
    full_link_slot_mask: chex.Array
    accepted_services: chex.Array
    accepted_bitrate: chex.Array
    total_bitrate: chex.Array
    list_of_requests: chex.Array

GNModelEnvParams

Bases: RSAEnvParams

Dataclass to hold environment state for GN model environments.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvParams(RSAEnvParams):
    """Dataclass to hold environment state for GN model environments.
    """
    ref_lambda: chex.Scalar = struct.field(pytree_node=False)
    max_spans: chex.Scalar = struct.field(pytree_node=False)
    max_span_length: chex.Scalar = struct.field(pytree_node=False)
    nonlinear_coeff: chex.Scalar = struct.field(pytree_node=False)
    raman_gain_slope: chex.Scalar = struct.field(pytree_node=False)
    attenuation: chex.Scalar = struct.field(pytree_node=False)
    attenuation_bar: chex.Scalar = struct.field(pytree_node=False)
    dispersion_coeff: chex.Scalar = struct.field(pytree_node=False)
    dispersion_slope: chex.Scalar = struct.field(pytree_node=False)
    transceiver_snr: chex.Array = struct.field(pytree_node=False)
    amplifier_noise_figure: chex.Array = struct.field(pytree_node=False)
    coherent: bool = struct.field(pytree_node=False)
    num_roadms: chex.Scalar = struct.field(pytree_node=False)
    roadm_loss: chex.Scalar = struct.field(pytree_node=False)
    num_spans: chex.Scalar = struct.field(pytree_node=False)
    launch_power_type: chex.Scalar = struct.field(pytree_node=False)
    snr_margin: chex.Scalar = struct.field(pytree_node=False)
    max_snr: chex.Scalar = struct.field(pytree_node=False)
    max_power: chex.Scalar = struct.field(pytree_node=False)
    min_power: chex.Scalar = struct.field(pytree_node=False)
    step_power: chex.Scalar = struct.field(pytree_node=False)
    last_fit: bool = struct.field(pytree_node=False)
    default_launch_power: chex.Scalar = struct.field(pytree_node=False)
    mod_format_correction: bool = struct.field(pytree_node=False)
    monitor_active_lightpaths: bool = struct.field(pytree_node=False)  # Monitor active lightpaths for throughput calculation
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)
    uniform_spans: bool = struct.field(pytree_node=False)

GNModelEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class GNModelEnvState(RSAEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    link_snr_array: chex.Array  # Available SNR on each link
    channel_centre_bw_array: chex.Array  # Channel centre bandwidth for each active connection
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots (used for lightpath SNR calculation)
    channel_power_array: chex.Array  # Channel power for each active connection
    channel_centre_bw_array_prev: chex.Array  # Channel centre bandwidth for each active connection in previous timestep
    path_index_array_prev: chex.Array  # Contains indices of lightpaths in use on slots in previous timestep
    channel_power_array_prev: chex.Array  # Channel power for each active connection in previous timestep
    launch_power_array: chex.Array  # Launch power array

LogEnvState

Dataclass to hold environment state for logging.

Parameters:

Name	Type	Description	Default
env_state	EnvState	Environment state	required
lengths	Scalar	Lengths	required
returns	Scalar	Returns	required
cum_returns	Scalar	Cumulative returns	required
episode_lengths	Scalar	Episode lengths	required
episode_returns	Scalar	Episode returns	required
accepted_services	Scalar	Accepted services	required
accepted_bitrate	Scalar	Accepted bitrate	required
done	Scalar	Done	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class LogEnvState:
    """Dataclass to hold environment state for logging.

    Args:
        env_state (EnvState): Environment state
        lengths (chex.Scalar): Lengths
        returns (chex.Scalar): Returns
        cum_returns (chex.Scalar): Cumulative returns
        episode_lengths (chex.Scalar): Episode lengths
        episode_returns (chex.Scalar): Episode returns
        accepted_services (chex.Scalar): Accepted services
        accepted_bitrate (chex.Scalar): Accepted bitrate
        done (chex.Scalar): Done
    """
    env_state: EnvState
    lengths: float
    returns: float
    cum_returns: float
    accepted_services: int
    accepted_bitrate: float
    total_bitrate: float
    utilisation: float
    done: bool

MultiBandRSAEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_start: chex.Scalar = struct.field(pytree_node=False)
    gap_width: chex.Scalar = struct.field(pytree_node=False)

MultiBandRSAEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class MultiBandRSAEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RMSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulations_array: chex.Array = struct.field(pytree_node=False)

RMSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RMSA with GN model.

Parameters:

Name	Type	Description	Default
link_snr_array	Array	Link SNR array	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RMSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RMSA with GN model.

    Args:
        link_snr_array (chex.Array): Link SNR array
    """
    modulation_format_index_array: chex.Array  # Modulation format index for each active connection
    modulation_format_index_array_prev: chex.Array  # Modulation format index for each active connection in previous timestep
    mod_format_mask: chex.Array  # Modulation format mask

RSAEnvParams

Bases: EnvParams

Dataclass to hold environment parameters for RSA.

Parameters:

Name	Type	Description	Default
num_nodes	Scalar	Number of nodes	required
num_links	Scalar	Number of links	required
link_resources	Scalar	Number of link resources	required
k_paths	Scalar	Number of paths	required
mean_service_holding_time	Scalar	Mean service holding time	required
load	Scalar	Load	required
arrival_rate	Scalar	Arrival rate	required
path_link_array	Array	Path link array	required
random_traffic	bool	Random traffic matrix for RSA on each reset (else uniform or custom)	required
max_slots	Scalar	Maximum number of slots	required
path_se_array	Array	Path spectral efficiency array	required
deterministic_requests	bool	If True, use deterministic requests	required
multiple_topologies	bool	If True, use multiple topologies	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvParams(EnvParams):
    """Dataclass to hold environment parameters for RSA.

    Args:
        num_nodes (chex.Scalar): Number of nodes
        num_links (chex.Scalar): Number of links
        link_resources (chex.Scalar): Number of link resources
        k_paths (chex.Scalar): Number of paths
        mean_service_holding_time (chex.Scalar): Mean service holding time
        load (chex.Scalar): Load
        arrival_rate (chex.Scalar): Arrival rate
        path_link_array (chex.Array): Path link array
        random_traffic (bool): Random traffic matrix for RSA on each reset (else uniform or custom)
        max_slots (chex.Scalar): Maximum number of slots
        path_se_array (chex.Array): Path spectral efficiency array
        deterministic_requests (bool): If True, use deterministic requests
        multiple_topologies (bool): If True, use multiple topologies
    """
    num_nodes: chex.Scalar = struct.field(pytree_node=False)
    num_links: chex.Scalar = struct.field(pytree_node=False)
    link_resources: chex.Scalar = struct.field(pytree_node=False)
    k_paths: chex.Scalar = struct.field(pytree_node=False)
    mean_service_holding_time: chex.Scalar = struct.field(pytree_node=False)
    load: chex.Scalar = struct.field(pytree_node=False)
    arrival_rate: chex.Scalar = struct.field(pytree_node=False)
    path_link_array: chex.Array = struct.field(pytree_node=False)
    random_traffic: bool = struct.field(pytree_node=False)
    max_slots: chex.Scalar = struct.field(pytree_node=False)
    path_se_array: chex.Array = struct.field(pytree_node=False)
    deterministic_requests: bool = struct.field(pytree_node=False)
    multiple_topologies: bool = struct.field(pytree_node=False)
    log_actions: bool = struct.field(pytree_node=False)
    disable_node_features: bool = struct.field(pytree_node=False)

RSAEnvState

Bases: EnvState

Dataclass to hold environment state for RSA.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
request_array	Array	Request array	required
link_slot_departure_array	Array	Link slot departure array	required
link_slot_mask	Array	Link slot mask	required
traffic_matrix	Array	Traffic matrix	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAEnvState(EnvState):
    """Dataclass to hold environment state for RSA.

    Args:
        link_slot_array (chex.Array): Link slot array
        request_array (chex.Array): Request array
        link_slot_departure_array (chex.Array): Link slot departure array
        link_slot_mask (chex.Array): Link slot mask
        traffic_matrix (chex.Array): Traffic matrix
    """
    link_slot_array: chex.Array
    request_array: chex.Array
    link_slot_departure_array: chex.Array
    link_slot_mask: chex.Array
    traffic_matrix: chex.Array

RSAGNModelEnvParams

Bases: GNModelEnvParams

Dataclass to hold environment params for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvParams(GNModelEnvParams):
    """Dataclass to hold environment params for RSA with GN model.
    """
    min_snr: chex.Scalar = struct.field(pytree_node=False)
    fec_threshold: chex.Scalar = struct.field(pytree_node=False)

RSAGNModelEnvState

Bases: GNModelEnvState

Dataclass to hold environment state for RSA with GN model.

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAGNModelEnvState(GNModelEnvState):
    """Dataclass to hold environment state for RSA with GN model.
    """
    active_lightpaths_array: chex.Array  # Active lightpath array. 1 x M array. Each value is a lightpath index. Used to calculate total throughput.
    active_lightpaths_array_departure: chex.Array  # Active lightpath array departure time.
    throughput: chex.Array  # Current network throughput

RSAMultibandEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for MultiBandRSA (RBSA).
    """
    gap_starts: chex.Array = struct.field(pytree_node=False)
    gap_widths: chex.Array = struct.field(pytree_node=False)

RSAMultibandEnvState

Bases: RSAEnvState

Dataclass to hold environment state for MultiBandRSA (RBSA).

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RSAMultibandEnvState(RSAEnvState):
    """Dataclass to hold environment state for MultiBandRSA (RBSA).
    """
    pass

RWALightpathReuseEnvState

Bases: RSAEnvState

Dataclass to hold environment state for RWA with lightpath reuse.

Parameters:

Name	Type	Description	Default
path_index_array	Array	Contains indices of lightpaths in use on slots	required
path_capacity_array	Array	Contains remaining capacity of each lightpath	required
link_capacity_array	Array	Contains remaining capacity of lightpath on each link-slot	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class RWALightpathReuseEnvState(RSAEnvState):
    """Dataclass to hold environment state for RWA with lightpath reuse.

    Args:
        path_index_array (chex.Array): Contains indices of lightpaths in use on slots
        path_capacity_array (chex.Array): Contains remaining capacity of each lightpath
        link_capacity_array (chex.Array): Contains remaining capacity of lightpath on each link-slot
    """
    path_index_array: chex.Array  # Contains indices of lightpaths in use on slots
    path_capacity_array: chex.Array  # Contains remaining capacity of each lightpath
    link_capacity_array: chex.Array  # Contains remaining capacity of lightpath on each link-slot
    time_since_last_departure: chex.Array  # Time since last departure

VONEEnvParams

Bases: RSAEnvParams

Dataclass to hold environment parameters for VONE.

Parameters:

Name	Type	Description	Default
node_resources	Scalar	Number of node resources	required
max_edges	Scalar	Maximum number of edges	required
min_node_resources	Scalar	Minimum number of node resources	required
max_node_resources	Scalar	Maximum number of node resources	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvParams(RSAEnvParams):
    """Dataclass to hold environment parameters for VONE.

    Args:
        node_resources (chex.Scalar): Number of node resources
        max_edges (chex.Scalar): Maximum number of edges
        min_node_resources (chex.Scalar): Minimum number of node resources
        max_node_resources (chex.Scalar): Maximum number of node resources
    """
    node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_edges: chex.Scalar = struct.field(pytree_node=False)
    min_node_resources: chex.Scalar = struct.field(pytree_node=False)
    max_node_resources: chex.Scalar = struct.field(pytree_node=False)

VONEEnvState

Bases: RSAEnvState

Dataclass to hold environment state for VONE.

Parameters:

Name	Type	Description	Default
node_capacity_array	Array	Node capacity array	required
node_resource_array	Array	Node resource array	required
node_departure_array	Array	Node departure array	required
action_counter	Array	Action counter	required
action_history	Array	Action history	required
node_mask_s	Array	Node mask for source node	required
node_mask_d	Array	Node mask for destination node	required
virtual_topology_patterns	Array	Virtual topology patterns	required
values_nodes	Array	Values for nodes	required

Source code in xlron/environments/dataclasses.py

@struct.dataclass
class VONEEnvState(RSAEnvState):
    """Dataclass to hold environment state for VONE.

    Args:
        node_capacity_array (chex.Array): Node capacity array
        node_resource_array (chex.Array): Node resource array
        node_departure_array (chex.Array): Node departure array
        action_counter (chex.Array): Action counter
        action_history (chex.Array): Action history
        node_mask_s (chex.Array): Node mask for source node
        node_mask_d (chex.Array): Node mask for destination node
        virtual_topology_patterns (chex.Array): Virtual topology patterns
        values_nodes (chex.Array): Values for nodes
    """
    node_capacity_array: chex.Array
    node_resource_array: chex.Array
    node_departure_array: chex.Array
    action_counter: chex.Array
    action_history: chex.Array
    node_mask_s: chex.Array
    node_mask_d: chex.Array
    virtual_topology_patterns: chex.Array
    values_nodes: chex.Array

aggregate_slots(full_mask, params)

Aggregate slot mask to reduce action space. Only used if the --aggregate_slots flag is set to > 1. Aggregated action is valid if there is one valid slot action within the aggregated action window.

Parameters:

Name	Type	Description	Default
full_mask	Array	slot mask	required
params	EnvParams	environment parameters	required

Returns:

Name	Type	Description
agg_mask	Array	aggregated slot mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def aggregate_slots(full_mask: chex.Array, params: EnvParams) -> chex.Array:
    """Aggregate slot mask to reduce action space. Only used if the --aggregate_slots flag is set to > 1.
    Aggregated action is valid if there is one valid slot action within the aggregated action window.

    Args:
        full_mask: slot mask
        params: environment parameters

    Returns:
        agg_mask: aggregated slot mask
    """

    num_actions = math.ceil(params.link_resources/params.aggregate_slots)
    agg_mask = jnp.zeros((params.k_paths, num_actions), dtype=LARGE_FLOAT_DTYPE)

    def get_max(i, mask_val):
        """Get maximum value of array slice of length aggregate_slots."""
        mask_slice = jax.lax.dynamic_slice(
                mask_val,
                (0, i * params.aggregate_slots,),
                (1,  params.aggregate_slots,),
            )
        max_slice = jnp.max(mask_slice).reshape(1, -1)
        return max_slice

    def update_window_max(i, val):
        """Update ith index 'agg_mask' with max of ith slice of length aggregate_slots from 'full_mask'.

        Args:
            i: increments as += aggregate_slots
            val: tuple of (agg_mask, path_mask, path_index).
        Returns:
            new_agg_mask: agg_mask is updated with max of path_mask for window size aggregate_slots
            mask: mask is unchanged
            path_index: path_index is unchanged
        """
        agg_mask = val[0]
        full_mask = val[1]
        path_index = val[2]
        new_agg_mask = jax.lax.dynamic_update_slice(
            agg_mask,
            get_max(i, full_mask),
            (path_index, i),
        )
        return new_agg_mask, full_mask, path_index

    def apply_to_path_mask(i, val):
        """
        Loop through each path for num_actions steps and get_window_max at each step.

        Args:
            i: path index
            val: tuple of (agg_mask, mask) where mask is original link-slot mask and agg_mask is resulting aggregated mask
        Returns:
            new_agg_mask: agg_mask is updated with aggregated path mask
            mask: mask is unchanged
        """
        val = (
            val[0],  # aggregated mask (to be updated)
            val[1][i].reshape(1, -1),  # mask for path i
            i  # path index
        )
        new_agg_mask = jax.lax.fori_loop(
            0,
            num_actions,
            update_window_max,
            val,
        )[0]
        return new_agg_mask, full_mask

    return jax.lax.fori_loop(
            0,
            params.k_paths,
            apply_to_path_mask,
            (agg_mask, full_mask),
        )

best_fit(state, params)

Best-Fit Spectrum Allocation. Returns the best fit slot for each path.

Source code in xlron/heuristics/heuristics.py

def best_fit(state: EnvState, params: EnvParams) -> chex.Array:
    """Best-Fit Spectrum Allocation. Returns the best fit slot for each path."""
    mask = get_action_mask(state, params)
    link_slot_array = jnp.where(state.link_slot_array < 0, 1., state.link_slot_array)
    nodes_sd, requested_bw = read_rsa_request(state.request_array)

    # We need to define a wrapper function in order to vmap with named arguments
    def _find_block_sizes(arr, starts_only=False, reverse=True):
        return jax.vmap(find_block_sizes, in_axes=(0, None, None))(arr, starts_only, reverse)

    block_sizes_right = _find_block_sizes(link_slot_array, starts_only=False, reverse=False)
    block_sizes_left = _find_block_sizes(link_slot_array, starts_only=False, reverse=True)
    block_sizes = jnp.maximum((block_sizes_left + block_sizes_right) - 1, 0)
    paths = get_paths(params, nodes_sd)
    se = get_paths_se(params, nodes_sd) if params.consider_modulation_format else jnp.ones((params.k_paths,))
    num_slots = jax.vmap(required_slots, in_axes=(None, 0, None, None))(requested_bw, se, params.slot_size, params.guardband)

    # Quantify how well the request fits within a free spectral block
    def get_bf_on_path(path, blocks, req_slots):
        fits = jax.vmap(lambda x: x - req_slots, in_axes=0)(blocks)
        fits = jnp.where(fits >= 0, fits, params.link_resources)
        path_fit = jnp.dot(path, fits) / jnp.sum(path)
        return path_fit
    fits_block = jax.vmap(lambda x, y, z: get_bf_on_path(x, y, z), in_axes=(0, None, 0))(paths, block_sizes, num_slots)

    # Quantity much of a gap there is between the assigned slots and the next occupied slots on the left
    def get_bf_on_path_left(path, blocks, req_slots):
        fits = jax.vmap(lambda x: x - req_slots, in_axes=0)(blocks)
        fits = jnp.where(fits >= 0, fits, params.link_resources)
        fits_shift = jax.lax.dynamic_slice(fits, (0, 1), (fits.shape[0], fits.shape[1]-1))
        fits_shift = jnp.concatenate((jnp.full((fits.shape[0], 1), params.link_resources), fits_shift), axis=1)
        fits = fits + 1/jnp.maximum(fits_shift, 1)
        path_fit = jnp.dot(path, fits) / jnp.sum(path)
        return path_fit
    fits_left = jax.vmap(lambda x, y, z: get_bf_on_path_left(x, y, z), in_axes=(0, None, 0))(paths, block_sizes_left, num_slots)

    # Quantity much of a gap there is between the assigned slots and the next occupied slots on the right
    def get_bf_on_path_right(path, blocks, req_slots):
        fits = jax.vmap(lambda x: x - req_slots, in_axes=0)(blocks)
        fits = jnp.where(fits >= 0, fits, params.link_resources)
        fits_shift = jax.lax.dynamic_slice(fits, (0, 0), (fits.shape[0], fits.shape[1] - 1))
        fits_shift = jnp.concatenate((fits_shift, jnp.full((fits.shape[0], 1), params.link_resources)), axis=1)
        fits = fits + 1/jnp.maximum(fits_shift, 1)
        path_fit = jnp.dot(path, fits) / jnp.sum(path)
        return path_fit
    fits_right = jax.vmap(lambda x, y, z: get_bf_on_path_right(x, y, z), in_axes=(0, None, 0))(paths, block_sizes_right, num_slots)

    # Sum the contribution to the overall quality of fit, and scale down the left/right contributions
    fits = jnp.sum(jnp.stack((fits_block, fits_left/params.link_resources, fits_right/params.link_resources), axis=0), axis=0)
    # Mask out occupied lightpaths (in case the quality of fit on some links is good enough to be considered, even if the overall path is invalid)
    fits = jnp.where(mask == 0, jnp.inf, fits)
    best_slots = jnp.argmin(fits, axis=1)
    best_fits = jnp.min(fits, axis=1)
    return best_slots, best_fits

bf_ksp(state, params)

Get the first available slot from the first k-shortest paths Method: Go through action mask and find the first available slot on all paths

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def bf_ksp(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the first available slot from the first k-shortest paths
    Method: Go through action mask and find the first available slot on all paths

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Action
    """
    best_slots, fitness = best_fit(state, params)
    # Chosen path is the one with the best fit
    path_index = jnp.argmin(fitness)
    slot_index = best_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

calculate_path_capacity(path_length, min_request=100, scale_factor=1.0, alpha=0.0002, NF=4.5, B=10000000000000.0, R_s=100000000000.0, beta_2=-2.17e-26, gamma=0.0012, L_s=100000.0, lambda0=1.55e-06)

From Nevin JOCN paper 2022: https://discovery.ucl.ac.uk/id/eprint/10175456/1/RL_JOCN_accepted.pdf

Source code in xlron/environments/env_funcs.py

def calculate_path_capacity(
        path_length,
        min_request=100,  # Minimum data rate request size
        scale_factor=1.0,  # Scale factor for link capacity
        alpha=0.2e-3, # Fibre attenuation coefficient
        NF=4.5,  # Amplifier noise figure
        B=10e12,  # Total modulated bandwidth
        R_s=100e9,  # Symbol rate
        beta_2=-21.7e-27,  # Dispersion parameter
        gamma=1.2e-3,  # Nonlinear coefficient
        L_s=100e3,  # Span length
        lambda0=1550e-9,  # Wavelength
):
    """From Nevin JOCN paper 2022: https://discovery.ucl.ac.uk/id/eprint/10175456/1/RL_JOCN_accepted.pdf"""
    alpha_lin = alpha / 4.343  # Linear attenuation coefficient
    N_spans = jnp.floor(path_length * 1e3 / L_s)  # Number of fibre spans on path
    L_eff = (1 - jnp.exp(-alpha_lin * L_s)) / alpha_lin  # Effective length of span in m
    sigma_2_ase = (jnp.exp(alpha_lin * L_s) - 1) * 10**(NF/10) * 6.626e-34 * 2.998e8 * R_s / lambda0  # ASE noise power
    span_NSR = jnp.cbrt(2 * sigma_2_ase**2 * alpha_lin * gamma**2 * L_eff**2 *
                        jnp.log(jnp.pi**2 * jnp.abs(beta_2) * B**2 / alpha_lin) / (jnp.pi * jnp.abs(beta_2) * R_s**2))  # Noise-to-signal ratio per span
    path_NSR = jnp.where(N_spans < 1, 1, N_spans) * span_NSR  # Noise-to-signal ratio per path
    path_capacity = 2 * R_s/1e9 * jnp.log2(1 + 1/path_NSR)  # Capacity of path in Gbps
    # Round link capacity down to nearest increment of minimum request size and apply scale factor
    path_capacity = jnp.floor(path_capacity * scale_factor / min_request) * min_request
    return path_capacity

calculate_path_stats(state, params, request)

For use in DeepRMSA agent observation space. Calculate: 1. Size of 1st suitable free spectrum block 2. Index of 1st suitable free spectrum block 3. Required slots on path 4. Avg. free block size 5. Free slots

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
stats	Array	Array of calculated path statistics

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def calculate_path_stats(state: EnvState, params: EnvParams, request: chex.Array) -> chex.Array:
    """For use in DeepRMSA agent observation space.
    Calculate:
    1. Size of 1st suitable free spectrum block
    2. Index of 1st suitable free spectrum block
    3. Required slots on path
    4. Avg. free block size
    5. Free slots

    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        stats: Array of calculated path statistics
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_val = jnp.zeros((params.k_paths, 5), dtype=LARGE_FLOAT_DTYPE)
    # TODO - check if the normalisation is useful
    def body_fun(i, val):
        link_resources = jnp.array(params.link_resources, dtype=LARGE_FLOAT_DTYPE)
        slot_size = jnp.array(params.slot_size, dtype=LARGE_FLOAT_DTYPE)
        slots = get_path_slots(state.link_slot_array, params, nodes_sd, i)
        se = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else jnp.array([1], dtype=SMALL_INT_DTYPE)
        req_slots = jnp.squeeze(required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband))
        req_slots_norm = req_slots * slot_size / jnp.max(params.values_bw.val)
        free_slots_norm = jnp.sum(jnp.where(slots == zero, one, zero), promote_integers=False) / link_resources
        block_sizes = find_block_sizes(slots)
        first_block_index = jnp.argmax(block_sizes >= req_slots)
        first_block_index_norm = first_block_index.astype(LARGE_FLOAT_DTYPE) / link_resources
        first_block_size_norm = jnp.squeeze(
            jax.lax.dynamic_slice(block_sizes, (first_block_index,), (1,))
        ) / req_slots.astype(LARGE_FLOAT_DTYPE)
        avg_block_size_norm = (jnp.sum(block_sizes) /
                               jnp.max(jnp.array([jnp.sum(find_block_starts(slots), promote_integers=False), 1])) /
                               req_slots)
        val = jax.lax.dynamic_update_slice(
            val,
            jnp.array([[
                first_block_size_norm,
                first_block_index_norm,
                req_slots_norm,
                avg_block_size_norm.astype(LARGE_FLOAT_DTYPE),
                free_slots_norm
            ]]),
            (i, 0)
        )  # N.B. that all values are normalised
        return val

    stats = jax.lax.fori_loop(
            0,
            params.k_paths,
            body_fun,
            init_val,
        )

    return stats

check_action_rmsa_gn_model(state, action, params)

Check if action is valid for RSA GN model Args: state (EnvState): Environment state params (EnvParams): Environment parameters action (chex.Array): Action array Returns: bool: True if action is invalid, False if action is valid

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def check_action_rmsa_gn_model(state: EnvState, action: Optional[chex.Array], params: EnvParams) -> bool:
    """Check if action is valid for RSA GN model
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        action (chex.Array): Action array
    Returns:
        bool: True if action is invalid, False if action is valid
    """
    # Check if action is valid
    # TODO - log failure reasons in info
    snr_sufficient_check = check_snr_sufficient(state, params)
    spectrum_reuse_check = check_no_spectrum_reuse(state.link_slot_array)
    # jax.debug.print("spectrum_reuse_check {}", spectrum_reuse_check, ordered=True)
    # jax.debug.print("snr_sufficient_check {}", snr_sufficient_check, ordered=True)
    return jnp.any(jnp.stack((
        spectrum_reuse_check,
        snr_sufficient_check,
    )))

check_action_rsa(state)

Check if action is valid. Combines checks for: - no spectrum reuse

Parameters:

Name	Type	Description	Default
state		current state	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_action_rsa(state):
    """Check if action is valid.
    Combines checks for:
    - no spectrum reuse

    Args:
        state: current state

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(jnp.stack((
        check_no_spectrum_reuse(state.link_slot_array),
    )))

check_action_rwalr(state, action, params)

Combines checks for: - no spectrum reuse - lightpath available and existing

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_action_rwalr(state: EnvState, action: chex.Array, params: EnvParams) -> bool:
    """Combines checks for:
    - no spectrum reuse
    - lightpath available and existing

    Args:
        state: Environment state

    Returns:
        bool: True if check failed, False if check passed

    """
    return jnp.any(jnp.stack((
        check_no_spectrum_reuse(state.link_slot_array),
        jnp.logical_not(check_lightpath_available_and_existing(state, params, action)[0]),
    )))

check_all_nodes_assigned(node_departure_array, total_requested_nodes)

Count negative values on each node (row) in node departure array, sum them, must equal total requested_nodes.

Parameters:

Name	Type	Description	Default
node_departure_array	Array	Node departure array (N x R) where N is number of nodes and R is number of resources	required
total_requested_nodes	int	Total requested nodes (int)	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_all_nodes_assigned(node_departure_array: chex.Array, total_requested_nodes: int) -> bool:
    """Count negative values on each node (row) in node departure array, sum them, must equal total requested_nodes.

    Args:
        node_departure_array: Node departure array (N x R) where N is number of nodes and R is number of resources
        total_requested_nodes: Total requested nodes (int)

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.sum(jnp.sum(jnp.where(node_departure_array < 0, 1, 0), axis=1)) != total_requested_nodes

check_lightpath_available_and_existing(state, params, action)

Check if lightpath is available and existing. Available means that the lightpath does not use slots occupied by a different lightpath. Existing means that the lightpath has already been established.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Name	Type	Description
lightpath_available_check	Tuple[Array, Array, Array, Array]	True if lightpath is available

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def check_lightpath_available_and_existing(state: EnvState, params: EnvParams, action: chex.Array) -> (
        Tuple)[chex.Array, chex.Array, chex.Array, chex.Array]:
    """Check if lightpath is available and existing.
    Available means that the lightpath does not use slots occupied by a different lightpath.
    Existing means that the lightpath has already been established.

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        lightpath_available_check: True if lightpath is available
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_index, initial_slot_index = process_path_action(state, params, action)
    path = get_paths(params, nodes_sd)[path_index]
    # Get unique lightpath index
    lightpath_index = get_lightpath_index(params, nodes_sd, path_index)
    # Get mask for slots that lightpath will occupy
    # negative numbers used so as not to conflict with lightpath indices
    new_lightpath_mask = vmap_set_path_links(
        jnp.full((params.num_links, 1), -2), path, 0, 1, -1
    )
    path_index_array = state.path_index_array[:, initial_slot_index].reshape(-1, 1)
    masked_path_index_array = jnp.where(
        new_lightpath_mask == -1, path_index_array, -2
    )
    lightpath_mask = jnp.where(
        path_index_array == lightpath_index, -1, -2
    )  # Allow current lightpath
    lightpath_existing_check = jnp.array_equal(lightpath_mask, new_lightpath_mask)  # True if all slots are same
    lightpath_mask = jnp.where(masked_path_index_array == -1, -1, lightpath_mask)  # Allow empty slots
    # True if all slots are same or empty
    lightpath_available_check = jnp.logical_or(
        jnp.array_equal(lightpath_mask, new_lightpath_mask), lightpath_existing_check
    )
    curr_lightpath_capacity = jnp.max(
        jnp.where(new_lightpath_mask == -1, state.link_capacity_array[:, initial_slot_index].reshape(-1, 1), 0)
    )
    return lightpath_available_check, lightpath_existing_check, curr_lightpath_capacity, lightpath_index

check_min_two_nodes_assigned(node_departure_array)

Count negative values on each node (row) in node departure array, sum them, must be 2 or greater. This check is important if e.g. an action contains 2 nodes the same therefore only assigns 1 node. Return False if check passed, True if check failed

Parameters:

Name	Type	Description	Default
node_departure_array	Array	Node departure array (N x R) where N is number of nodes and R is number of resources	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_min_two_nodes_assigned(node_departure_array: chex.Array):
    """Count negative values on each node (row) in node departure array, sum them, must be 2 or greater.
    This check is important if e.g. an action contains 2 nodes the same therefore only assigns 1 node.
    Return False if check passed, True if check failed

    Args:
        node_departure_array: Node departure array (N x R) where N is number of nodes and R is number of resources

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.sum(jnp.sum(jnp.where(node_departure_array < 0, 1, 0), axis=1)) <= 1

check_no_spectrum_reuse(link_slot_array)

slot-=1 when used, should be zero when unoccupied, so check if any < -1 in slot array.

Parameters:

Name	Type	Description	Default
link_slot_array		Link slot array (L x S) where L is number of links and S is number of slots	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_no_spectrum_reuse(link_slot_array):
    """slot-=1 when used, should be zero when unoccupied, so check if any < -1 in slot array.

    Args:
        link_slot_array: Link slot array (L x S) where L is number of links and S is number of slots

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(link_slot_array < -1)

check_node_capacities(capacity_array)

Sum selected nodes array and check less than node resources.

Parameters:

Name	Type	Description	Default
capacity_array	Array	Node capacity array (N x 1) where N is number of nodes	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_node_capacities(capacity_array: chex.Array) -> bool:
    """Sum selected nodes array and check less than node resources.

    Args:
        capacity_array: Node capacity array (N x 1) where N is number of nodes

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(capacity_array < 0)

check_snr_sufficient(state, params)

Check if SNR is sufficient for all connections Args: state (EnvState): Environment state params (EnvParams): Environment parameters Returns: jnp.array: 1 if SNR is sufficient for connection else 0

Source code in xlron/environments/env_funcs.py

def check_snr_sufficient(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> chex.Array:
    """Check if SNR is sufficient for all connections
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: 1 if SNR is sufficient for connection else 0
    """
    # TODO - this check needs to be faster!
    required_snr_array = get_required_snr_se_kurtosis_array(state.modulation_format_index_array, 2, params)
    # Transform lightpath index array by getting lightpath value, getting path-link array, and summing inverse link SNRs
    lightpath_snr_array = get_lightpath_snr(state, params)
    check_snr_sufficient = jnp.where(lightpath_snr_array >= required_snr_array, 0, 1)
    # jax.debug.print("check_snr_sufficient {}", check_snr_sufficient, ordered=True)
    # jax.debug.print("required_snr_array {}", required_snr_array, ordered=True)
    # jax.debug.print("lightpath_snr_array {}", lightpath_snr_array, ordered=True)
    # jax.debug.print("state.modulation_format_index_array {}", state.modulation_format_index_array, ordered=True)
    # jax.debug.print("state.channel_centre_bw_array {}", state.channel_centre_bw_array, ordered=True)
    # jax.debug.print("state.channel_power_array {}", state.channel_power_array, ordered=True)
    return jnp.any(check_snr_sufficient)

check_topology(action_history, topology_pattern)

Check that each unique virtual node (as indicated by topology pattern) is assigned to a consistent physical node i.e. start and end node of ring is same physical node. Method: For each node index in topology pattern, mask action history with that index, then find max value in masked array. If max value is not the same for all values for that virtual node in action history, then return 1, else 0. Array should be all zeroes at the end, so do an any() check on that. e.g. virtual topology pattern = [2,1,3,1,4,1,2] 3 node ring action history = [0,34,4,0,3,1,0] meaning v node "2" goes p node 0, v node "3" goes p node 4, v node "4" goes p node 3 The numbers in-between relate to the slot action. If any value in the array is 1, a virtual node is assigned to multiple different physical nodes. Need to check from both perspectives: 1. For each virtual node, check that all physical nodes are the same 2. For each physical node, check that all virtual nodes are the same

Parameters:

Name	Type	Description	Default
action_history		Action history	required
topology_pattern		Topology pattern	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_topology(action_history, topology_pattern):
    """Check that each unique virtual node (as indicated by topology pattern) is assigned to a consistent physical node
    i.e. start and end node of ring is same physical node.
    Method:
    For each node index in topology pattern, mask action history with that index, then find max value in masked array.
    If max value is not the same for all values for that virtual node in action history, then return 1, else 0.
    Array should be all zeroes at the end, so do an any() check on that.
    e.g. virtual topology pattern = [2,1,3,1,4,1,2]  3 node ring
    action history = [0,34,4,0,3,1,0]
    meaning v node "2" goes p node 0, v node "3" goes p node 4, v node "4" goes p node 3
    The numbers in-between relate to the slot action.
    If any value in the array is 1, a virtual node is assigned to multiple different physical nodes.
    Need to check from both perspectives:
    1. For each virtual node, check that all physical nodes are the same
    2. For each physical node, check that all virtual nodes are the same

    Args:
        action_history: Action history
        topology_pattern: Topology pattern

    Returns:
        bool: True if check failed, False if check passed
    """
    def loop_func_virtual(i, val):
        # Get indices of physical node in action history that correspond to virtual node i
        masked_val = jnp.where(i == topology_pattern, val, -1)
        # Get maximum value at those indices (should all be same)
        max_node = jnp.max(masked_val)
        # For relevant indices, if max value then return 0 else 1
        val = jnp.where(masked_val != -1, masked_val != max_node, val)
        return val
    def loop_func_physical(i, val):
        # Get indices of virtual nodes in topology pattern that correspond to physical node i
        masked_val = jnp.where(i == action_history, val, -1)
        # Get maximum value at those indices (should all be same)
        max_node = jnp.max(masked_val)
        # For relevant indices, if max value then return 0 else 1
        val = jnp.where(masked_val != -1, masked_val != max_node, val)
        return val
    topology_pattern = topology_pattern[::2]  # Only look at node indices, not slot actions
    action_history = action_history[::2]
    check_virtual = jax.lax.fori_loop(jnp.min(topology_pattern), jnp.max(topology_pattern)+1, loop_func_virtual, action_history)
    check_physical = jax.lax.fori_loop(jnp.min(action_history), jnp.max(action_history)+1, loop_func_physical, topology_pattern)
    check = jnp.concatenate((check_virtual, check_physical))
    return jnp.any(check)

check_unique_nodes(node_departure_array)

Count negative values on each node (row) in node departure array, must not exceed 1.

Parameters:

Name	Type	Description	Default
node_departure_array	Array	Node departure array (N x R) where N is number of nodes and R is number of resources	required

Returns:

Name	Type	Description
bool	bool	True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

@jax.jit
def check_unique_nodes(node_departure_array: chex.Array) -> bool:
    """Count negative values on each node (row) in node departure array, must not exceed 1.

    Args:
        node_departure_array: Node departure array (N x R) where N is number of nodes and R is number of resources

    Returns:
        bool: True if check failed, False if check passed
    """
    return jnp.any(jnp.sum(jnp.where(node_departure_array < zero, one, zero), axis=1, promote_integers=False) > one)

check_vone_action(state, remaining_actions, total_requested_nodes)

Check if action is valid. Combines checks for: - sufficient node capacities - unique nodes assigned - minimum two nodes assigned - all requested nodes assigned - correct topology pattern - no spectrum reuse

Parameters:

Name	Type	Description	Default
state		current state	required
remaining_actions		remaining actions	required
total_requested_nodes		total requested nodes	required

Returns:

Name	Type	Description
bool		True if check failed, False if check passed

Source code in xlron/environments/env_funcs.py

def check_vone_action(state, remaining_actions, total_requested_nodes):
    """Check if action is valid.
    Combines checks for:
    - sufficient node capacities
    - unique nodes assigned
    - minimum two nodes assigned
    - all requested nodes assigned
    - correct topology pattern
    - no spectrum reuse

    Args:
        state: current state
        remaining_actions: remaining actions
        total_requested_nodes: total requested nodes

    Returns:
        bool: True if check failed, False if check passed
    """
    checks = jnp.stack((
        check_node_capacities(state.node_capacity_array),
        check_unique_nodes(state.node_departure_array),
        # TODO (VONE) - Remove two nodes check if impairs performance
        #  (check_all_nodes_assigned is sufficient but fails after last action of request instead of earlier)
        check_min_two_nodes_assigned(state.node_departure_array),
        jax.lax.cond(
            jnp.equal(remaining_actions, jnp.array(1)),
            lambda x: check_all_nodes_assigned(*x),
            lambda x: jnp.array(False),
            (state.node_departure_array, total_requested_nodes)
        ),
        jax.lax.cond(
            jnp.equal(remaining_actions, jnp.array(1)),
            lambda x: check_topology(*x),
            lambda x: jnp.array(False),
            (state.action_history, state.request_array[1])
        ),
        check_no_spectrum_reuse(state.link_slot_array),
    ))
    #jax.debug.print("Checks: {}", checks, ordered=True)
    return jnp.any(checks)

convert_node_probs_to_traffic_matrix(node_probs)

Convert list of node probabilities to symmetric traffic matrix.

Parameters:

Name	Type	Description	Default
node_probs	list	node probabilities	required

Returns:

Name	Type	Description
traffic_matrix	Array	traffic matrix

Source code in xlron/environments/env_funcs.py

def convert_node_probs_to_traffic_matrix(node_probs: list) -> chex.Array:
    """Convert list of node probabilities to symmetric traffic matrix.

    Args:
        node_probs: node probabilities

    Returns:
        traffic_matrix: traffic matrix
    """
    matrix = jnp.outer(node_probs, node_probs).astype(SMALL_FLOAT_DTYPE)
    # Set lead diagonal to zero
    matrix = jnp.where(jnp.eye(matrix.shape[0]) == 1, 0, matrix)
    matrix = normalise_traffic_matrix(matrix)
    return matrix

create_run_name(config)

Create name for run based on config flags

Source code in xlron/environments/env_funcs.py

def create_run_name(config: Union[box.Box, dict]) -> str:
    """Create name for run based on config flags"""
    env_type = config["env_type"]
    topology = config["topology_name"]
    slots = config["link_resources"]
    gnn = "_GNN" if config["USE_GNN"] else ""
    incremental = "_INC" if config["incremental_loading"] else ""
    run_name = f"{env_type}_{topology}_{slots}{gnn}{incremental}".upper()
    if config["EVAL_HEURISTIC"]:
        run_name += f"_{config['path_heuristic']}"
        if env_type.lower() == "vone":
            run_name += f"_{config['node_heuristic']}"
    elif config["EVAL_MODEL"]:
        run_name += f"_EVAL"
    return run_name

decrement_action_counter(state)

Decrement action counter in-place. Used in VONE environments.

Source code in xlron/environments/env_funcs.py

def decrement_action_counter(state):
    """Decrement action counter in-place. Used in VONE environments."""
    state.action_counter.at[-1].add(-1)
    return state

ff_ksp(state, params)

Get the first available slot from all paths Method: Go through action mask and find the first available slot on all paths

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def ff_ksp(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the first available slot from all paths
    Method: Go through action mask and find the first available slot on all paths

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Action
    """
    first_slots = first_fit(state, params)
    # Chosen path is the one with the lowest index of first available slot
    path_index = jnp.argmin(first_slots)
    slot_index = first_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

finalise_action_rsa(state, params)

Turn departure times positive.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def finalise_action_rsa(state: EnvState, params: Optional[EnvParams]):
    """Turn departure times positive.

    Args:
        state: current state

    Returns:
        state: updated state
    """
    _, requested_datarate = read_rsa_request(state.request_array)
    state = state.replace(
        link_slot_departure_array=make_positive(state.link_slot_departure_array),
        accepted_services=state.accepted_services + 1,
        accepted_bitrate=state.accepted_bitrate + requested_datarate[0],
        total_bitrate=state.total_bitrate + requested_datarate[0]
    )
    return state

finalise_action_rwalr(state, params)

Turn departure times positive.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def finalise_action_rwalr(state: EnvState, params: Optional[EnvParams]):
    """Turn departure times positive.

    Args:
        state: current state

    Returns:
        state: updated state
    """
    _, requested_datarate = read_rsa_request(state.request_array)
    state = state.replace(
        link_slot_departure_array=make_positive(state.link_slot_departure_array),
        accepted_services=state.accepted_services + 1,
        accepted_bitrate=state.accepted_bitrate + requested_datarate[0],
        total_bitrate=state.total_bitrate + requested_datarate[0]
    )
    return state

finalise_vone_action(state)

Turn departure times positive.

Parameters:

Name	Type	Description	Default
state		current state	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def finalise_vone_action(state):
    """Turn departure times positive.

    Args:
        state: current state

    Returns:
        state: updated state
    """
    state = state.replace(
        node_departure_array=make_positive(state.node_departure_array),
        link_slot_departure_array=make_positive(state.link_slot_departure_array),
        accepted_services=state.accepted_services + 1,
        accepted_bitrate=state.accepted_bitrate  # TODO - get sum of bitrates for requested links
    )
    return state

first_fit(state, params)

First-Fit Spectrum Allocation. Returns the first fit slot for each path.

Source code in xlron/heuristics/heuristics.py

def first_fit(state: EnvState, params: EnvParams) -> chex.Array:
    """First-Fit Spectrum Allocation. Returns the first fit slot for each path."""
    mask = get_action_mask(state, params)
    # Add a column of ones to the mask to make sure that occupied paths have non-zero index in "first_slots"
    mask = jnp.concatenate((mask, jnp.full((mask.shape[0], 1), 1)), axis=1)
    # Get index of first available slots for each path
    first_slots = jnp.argmax(mask, axis=1)
    return first_slots

format_vone_slot_request(state, action)

Format slot request for VONE action into format (source-node, slot, destination-node).

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required
action	Array	action to format	required

Returns:

Type	Description
Array	chex.Array: formatted request

Source code in xlron/environments/env_funcs.py

def format_vone_slot_request(state: EnvState, action: chex.Array) -> chex.Array:
    """Format slot request for VONE action into format (source-node, slot, destination-node).

    Args:
        state: current state
        action: action to format

    Returns:
        chex.Array: formatted request
    """
    remaining_actions = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.action_counter, 2, 1))
    full_request = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.request_array, 0, 1))
    unformatted_request = jax.lax.dynamic_slice_in_dim(full_request, (remaining_actions - 1) * 2, 3)
    node_s = jax.lax.dynamic_slice_in_dim(action, 0, 1)
    requested_slots = jax.lax.dynamic_slice_in_dim(unformatted_request, 1, 1)
    node_d = jax.lax.dynamic_slice_in_dim(action, 2, 1)
    formatted_request = jnp.concatenate((node_s, requested_slots, node_d))
    return formatted_request

generate_arrival_holding_times(key, params)

Generate arrival and holding times based on Poisson distributed events. To understand how sampling from e^-x can be transformed to sample from lambdae^-(x/lambda) see: https://en.wikipedia.org/wiki/Inverse_transform_sampling#Examples Basically, inverse transform sampling is used to sample from a distribution with CDF F(x). The CDF of the exponential distribution (lambdae^-{lambdax}) is F(x) = 1 - e^-{lambdax}. Therefore, the inverse CDF is x = -ln(1-u)/lambda, where u is sample from uniform distribution. Therefore, we need to divide jax.random.exponential() by lambda in order to scale the standard exponential CDF. Experimental histograms of this method compared to random.expovariate() in Python's random library show that the two methods are equivalent. Also see: https://numpy.org/doc/stable/reference/random/generated/numpy.random.exponential.html https://jax.readthedocs.io/en/latest/_autosummary/jax.random.exponential.html

Parameters:

Name	Type	Description	Default
key		PRNG key	required
params		Environment parameters	required

Returns:

Name	Type	Description
arrival_time		Arrival time
holding_time		Holding time

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def generate_arrival_holding_times(key, params):
    """
    Generate arrival and holding times based on Poisson distributed events.
    To understand how sampling from e^-x can be transformed to sample from lambda*e^-(x/lambda) see:
    https://en.wikipedia.org/wiki/Inverse_transform_sampling#Examples
    Basically, inverse transform sampling is used to sample from a distribution with CDF F(x).
    The CDF of the exponential distribution (lambda*e^-{lambda*x}) is F(x) = 1 - e^-{lambda*x}.
    Therefore, the inverse CDF is x = -ln(1-u)/lambda, where u is sample from uniform distribution.
    Therefore, we need to divide jax.random.exponential() by lambda in order to scale the standard exponential CDF.
    Experimental histograms of this method compared to random.expovariate() in Python's random library show that
    the two methods are equivalent.
    Also see: https://numpy.org/doc/stable/reference/random/generated/numpy.random.exponential.html
    https://jax.readthedocs.io/en/latest/_autosummary/jax.random.exponential.html

    Args:
        key: PRNG key
        params: Environment parameters

    Returns:
        arrival_time: Arrival time
        holding_time: Holding time
    """
    key_arrival, key_holding = jax.random.split(key, 2)
    arrival_time = jax.random.exponential(key_arrival, shape=(1,), dtype=SMALL_FLOAT_DTYPE) \
                   / params.arrival_rate  # Divide because it is rate (lambda)
    if params.truncate_holding_time:
        # For DeepRMSA, need to generate holding times that are less than 2*mean_service_holding_time
        key_holding = jax.random.split(key, 5)
        holding_times = jax.vmap(lambda x: jax.random.exponential(x, shape=(1,)) \
                                * params.mean_service_holding_time)(key_holding)
        holding_times = jnp.where(holding_times < 2*params.mean_service_holding_time, holding_times, zero)
        # Get first non-zero value in holding_times
        non_zero_index = jnp.nonzero(holding_times, size=1)[0][0]
        holding_time = jax.lax.dynamic_slice(jnp.squeeze(holding_times), (non_zero_index,), (1,))
    else:
        holding_time = jax.random.exponential(key_holding, shape=(1,), dtype=SMALL_FLOAT_DTYPE) \
                       * params.mean_service_holding_time  # Multiply because it is mean (1/lambda)
    return arrival_time, holding_time

generate_vone_request(key, state, params)

Generate a new request for the VONE environment. The request has two rows. The first row shows the node and slot values. The first three elements of the second row show the number of unique nodes, the total number of steps, and the remaining steps. These first three elements comprise the action counter. The remaining elements of the second row show the virtual topology pattern, i.e. the connectivity of the virtual topology.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def generate_vone_request(key: chex.PRNGKey, state: EnvState, params: EnvParams):
    """Generate a new request for the VONE environment.
    The request has two rows. The first row shows the node and slot values.
    The first three elements of the second row show the number of unique nodes, the total number of steps, and the remaining steps.
    These first three elements comprise the action counter.
    The remaining elements of the second row show the virtual topology pattern, i.e. the connectivity of the virtual topology.
    """
    shape = params.max_edges*2+1  # shape of request array
    key_topology, key_node, key_slot, key_times = jax.random.split(key, 4)
    # Randomly select topology, node resources, slot resources
    pattern = jax.random.choice(key_topology, state.virtual_topology_patterns)
    action_counter = jax.lax.dynamic_slice(pattern, (0,), (3,))
    topology_pattern = jax.lax.dynamic_slice(pattern, (3,), (pattern.shape[0]-3,))
    selected_node_values = jax.random.choice(key_node, state.values_nodes, shape=(shape,))
    selected_bw_values = jax.random.choice(key_slot, params.values_bw.val, shape=(shape,))
    # Create a mask for odd and even indices
    mask = jnp.tile(jnp.array([0, 1]), (shape+1) // 2)[:shape]
    # Vectorized conditional replacement using mask
    first_row = jnp.where(mask, selected_bw_values, selected_node_values)
    # Make sure node request values are consistent for same virtual nodes
    first_row = jax.lax.fori_loop(
        2,  # Lowest node index in virtual topology requests is 2
        shape,  # Highest possible node index in virtual topology requests is shape-1
        lambda i, x: jnp.where(topology_pattern == i, selected_node_values[i], x),
        first_row
    )
    # Mask out unused part of request array
    first_row = jnp.where(topology_pattern == 0, 0, first_row)
    # Set times
    arrival_time, holding_time = generate_arrival_holding_times(key, params)
    state = state.replace(
        holding_time=holding_time,
        current_time=state.current_time + arrival_time,
        action_counter=action_counter,
        request_array=jnp.vstack((first_row, topology_pattern)),
        action_history=init_action_history(params),
        total_requests=state.total_requests + 1
    )
    state = remove_expired_node_requests(state, params) if not params.incremental_loading else state
    state = remove_expired_services_rsa(state, params) if not params.incremental_loading else state
    return state

get_action_mask(state, params)

N.B. The mask must already be present in the state!

Source code in xlron/heuristics/heuristics.py

def get_action_mask(state: EnvState, params: EnvParams) -> chex.Array:
    """N.B. The mask must already be present in the state!"""
    mask = state.link_slot_mask if params.__class__.__name__ != "DeepRMSAEnvParams" else (
        mask_slots(state, params, state.request_array).link_slot_mask
    )
    mask = jnp.reshape(mask, (params.k_paths, -1))
    return mask

get_best_modulation_format(state, path, initial_slot_index, launch_power, params)

Get best modulation format for lightpath. "Best" is the highest order that has SNR requirements below available. Try each modulation format, calculate SNR for each, then return the highest order possible. Args: state (EnvState): Environment state path (chex.Array): Path array initial_slot_index (int): Initial slot index params (EnvParams): Environment parameters Returns: jnp.array: Acceptable modulation format indices

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(3,))
def get_best_modulation_format(state: EnvState, path: chex.Array, initial_slot_index: int, launch_power: chex.Array, params: EnvParams) -> chex.Array:
    """Get best modulation format for lightpath. "Best" is the highest order that has SNR requirements below available.
    Try each modulation format, calculate SNR for each, then return the highest order possible.
    Args:
        state (EnvState): Environment state
        path (chex.Array): Path array
        initial_slot_index (int): Initial slot index
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Acceptable modulation format indices
    """
    _, requested_datarate = read_rsa_request(state.request_array)
    mod_format_count = params.modulations_array.val.shape[0]
    acceptable_mod_format_indices = jnp.full((mod_format_count,), -2)

    def acceptable_modulation_format(i, acceptable_format_indices):
        req_snr = params.modulations_array.val[i][2] + params.snr_margin
        se = params.modulations_array.val[i][1]
        req_slots = required_slots(requested_datarate, se, params.slot_size, params.guardband)
        # TODO - need to check we don't overwrite values in already occupied slots
        # Possible approaches:
        # Check slot occupancy? Probably would need to iterate through for num_slots, but that's an issue
        # What about we allocate and then fix up later, e.g. could it be possible to just add the modulation format on top without
        # check sum of path links prior to assigning?
        #
        new_state = state.replace(
            channel_power_array=vmap_set_path_links(
                state.channel_power_array, path, initial_slot_index, req_slots, launch_power),
            channel_centre_bw_array=vmap_set_path_links(
                state.channel_centre_bw_array, path, initial_slot_index, req_slots, params.slot_size)
        )
        snr_value = get_minimum_snr_of_channels_on_path(new_state, path, initial_slot_index, req_slots, params)
        # jax.debug.print("snr_value {}", snr_value, ordered=True)
        # jax.debug.print("req_snr {}", req_snr, ordered=True)
        acceptable_format_index = jnp.where(snr_value >= req_snr, i, -1).reshape((1,))
        acceptable_format_indices = jax.lax.dynamic_update_slice(acceptable_format_indices, acceptable_format_index, (i,))
        # jax.debug.print("acceptable_format_indices {}", acceptable_format_indices, ordered=True)
        return acceptable_format_indices

    acceptable_mod_format_indices = jax.lax.fori_loop(
        0,
        mod_format_count,
        acceptable_modulation_format,
        acceptable_mod_format_indices
    )
    return acceptable_mod_format_indices

get_best_modulation_format_simple(state, path, initial_slot_index, params)

Get modulation format for lightpath. Assume worst case (least Gaussian) modulation format when calculating SNR. Args: state (EnvState): Environment state path (chex.Array): Path array initial_slot_index (int): Initial slot index params (EnvParams): Environment parameters Returns: jnp.array: Acceptable modulation format indices

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(3,))
def get_best_modulation_format_simple(
        state: RSAGNModelEnvState, path: chex.Array, initial_slot_index: int, params: RSAGNModelEnvParams
) -> chex.Array:
    """Get modulation format for lightpath.
    Assume worst case (least Gaussian) modulation format when calculating SNR.
    Args:
        state (EnvState): Environment state
        path (chex.Array): Path array
        initial_slot_index (int): Initial slot index
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Acceptable modulation format indices
    """
    link_snr_array = get_snr_link_array(state, params)
    snr_value = get_snr_for_path(path, link_snr_array, params)[initial_slot_index] - params.snr_margin  # Margin
    mod_format_count = params.modulations_array.val.shape[0]
    acceptable_mod_format_indices = jnp.arange(mod_format_count)
    req_snr = params.modulations_array.val[:, 2] + params.snr_margin
    acceptable_mod_format_indices = jnp.where(snr_value >= req_snr,
                                              acceptable_mod_format_indices,
                                              jnp.full((mod_format_count,), -2))
    return acceptable_mod_format_indices

get_centre_frequency(initial_slot_index, num_slots, params)

Get centre frequency for new lightpath

Parameters:

Name	Type	Description	Default
initial_slot_index	Array	Centre frequency of first slot	required
num_slots	float	Number of slots	required
params	RSAGNModelEnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Centre frequency for new lightpath

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def get_centre_frequency(initial_slot_index: int, num_slots: int, params: RSAGNModelEnvParams) -> chex.Array:
    """Get centre frequency for new lightpath

    Args:
        initial_slot_index (chex.Array): Centre frequency of first slot
        num_slots (float): Number of slots
        params (RSAGNModelEnvParams): Environment parameters

    Returns:
        chex.Array: Centre frequency for new lightpath
    """
    slot_centres = (jnp.arange(params.link_resources) - (params.link_resources - 1) / 2) * params.slot_size
    return slot_centres[initial_slot_index] + ((params.slot_size * (num_slots - 1)) / 2)

get_edge_disjoint_paths(graph)

Get edge disjoint paths between all nodes in graph.

Parameters:

Name	Type	Description	Default
graph	Graph	graph	required

Returns:

Name	Type	Description
dict	dict	edge disjoint paths (path is list of edges)

Source code in xlron/environments/env_funcs.py

def get_edge_disjoint_paths(graph: nx.Graph) -> dict:
    """Get edge disjoint paths between all nodes in graph.

    Args:
        graph: graph

    Returns:
        dict: edge disjoint paths (path is list of edges)
    """
    result = {n: {} for n in graph}
    for n1, n2 in itertools.combinations(graph, 2):
        # Sort by number of links in path
        # TODO - sort by path length
        result[n1][n2] = sorted(list(nx.edge_disjoint_paths(graph, n1, n2)), key=len)
        result[n2][n1] = sorted(list(nx.edge_disjoint_paths(graph, n2, n1)), key=len)
    return result

get_launch_power(state, path_action, power_action, params)

Get launch power for new lightpath. N.B. launch power is specified in dBm but is converted to linear units when stored in channel_power_array. This func returns linear units (mW). Path action is used to determine the launch power in the case of tabular launch power type. Power action is used to determine the launch power in the case of RL launch power type. During masking, power action is set as state.launch_power_array[0], which is set by the RL agent. Args: state (EnvState): Environment state path_action (chex.Array): Action specifying path index (0 to k_paths-1) power_action (chex.Array): Action specifying launch power in dBm params (EnvParams): Environment parameters Returns: chex.Array: Launch power for new lightpath

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_launch_power(state: EnvState, path_action: chex.Array, power_action: chex.Array, params: EnvParams) -> chex.Array:
    """Get launch power for new lightpath. N.B. launch power is specified in dBm but is converted to linear units
    when stored in channel_power_array. This func returns linear units (mW).
    Path action is used to determine the launch power in the case of tabular launch power type.
    Power action is used to determine the launch power in the case of RL launch power type. During masking,
    power action is set as state.launch_power_array[0], which is set by the RL agent.
    Args:
        state (EnvState): Environment state
        path_action (chex.Array): Action specifying path index (0 to k_paths-1)
        power_action (chex.Array): Action specifying launch power in dBm
        params (EnvParams): Environment parameters
    Returns:
        chex.Array: Launch power for new lightpath
    """
    k_path_index, _ = process_path_action(state, params, path_action)
    if params.launch_power_type == 1:  # Fixed
        return state.launch_power_array[0]
    elif params.launch_power_type == 2:  # Tabular (one row per path)
        nodes_sd, requested_datarate = read_rsa_request(state.request_array)
        source, dest = nodes_sd
        i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
        return state.launch_power_array[i+k_path_index]
    elif params.launch_power_type == 3:  # RL
        return power_action
    elif params.launch_power_type == 4:  # Scaled
        nodes_sd, requested_datarate = read_rsa_request(state.request_array)
        source, dest = nodes_sd
        i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
        # Get path length
        link_length_array = jnp.sum(params.link_length_array.val, axis=1, promote_integers=False)
        path_length = jnp.sum(link_length_array[i+k_path_index], promote_integers=False)
        path_link_array = jnp.unpackbits(params.path_link_array.val)[:, params.num_links] if params.pack_path_bits \
            else params.path_link_array.val
        maximum_path_length = jnp.max(jnp.dot(path_link_array, params.link_length_array.val))
        return state.launch_power_array[0] * (path_length / maximum_path_length)
    else:
        raise ValueError("Invalid launch power type. Check params.launch_power_type")

get_lightpath_snr(state, params)

Get SNR for each link on path. N.B. that in most cases it is more efficient to calculate the SNR for every possible path, rather than a slot-by-slot basis. But in some cases slot-by-slot is better i.e. when kN(N-1)/2 > LS Args: state (RSAGNModelEnvState): Environment state params (RSAGNModelEnvParams): Environment parameters

Returns:

Type	Description
Array	chex.array: SNR for each link on path

Source code in xlron/environments/env_funcs.py

def get_lightpath_snr(state: RSAGNModelEnvParams, params: RSAGNModelEnvParams) -> chex.Array:
    """Get SNR for each link on path.
    N.B. that in most cases it is more efficient to calculate the SNR for every possible path, rather than a slot-by-slot basis.
    But in some cases slot-by-slot is better i.e. when k*N(N-1)/2 > L*S
    Args:
        state (RSAGNModelEnvState): Environment state
        params (RSAGNModelEnvParams): Environment parameters

    Returns:
        chex.array: SNR for each link on path
    """
    # Get the SNR for the channel that the path occupies
    path_snr_array = jax.vmap(get_snr_for_path, in_axes=(0, None, None))(params.path_link_array.val, state.link_snr_array, params)
    # Where value in path_index_array matches index of path_snr_array, substitute in SNR value
    slot_indices = jnp.arange(params.link_resources)
    lightpath_snr_array = jax.vmap(jax.vmap(lambda x, si: path_snr_array[x][si], in_axes=(0, 0)), in_axes=(0, None))(state.path_index_array, slot_indices)
    return lightpath_snr_array

get_link_weights(state, params)

Get link weights based on occupancy for use in congestion-aware routing heuristics.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
	chex.Array: Link weights

Source code in xlron/heuristics/heuristics.py

def get_link_weights(state: EnvState, params: EnvParams):
    """Get link weights based on occupancy for use in congestion-aware routing heuristics.

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Link weights
    """
    if params.__class__.__name__ != "RWALightpathReuseEnvParams":
        link_occupancy = jnp.count_nonzero(state.link_slot_array, axis=1)
    else:
        initial_path_capacity = init_path_capacity_array(
            params.link_length_array.val, params.path_link_array.val, scale_factor=1.0
        )
        initial_path_capacity = jnp.squeeze(jax.vmap(lambda x: initial_path_capacity[x])(state.path_index_array))
        utilisation = jnp.where(initial_path_capacity - state.link_capacity_array < 0, 0,
                                initial_path_capacity - state.link_capacity_array) / initial_path_capacity
        link_occupancy = jnp.sum(utilisation, axis=1)
    link_weights = jnp.multiply(params.link_length_array.val.T, (1 / (1 - link_occupancy / (params.link_resources + 1))))[0]
    return link_weights

get_minimum_snr_of_channels_on_path(state, path, slot_index, req_slots, params)

Get the minimum value of the SNR on newly assigned channels. N.B. this requires the link_snr_array to have already been calculated and present in state.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def get_minimum_snr_of_channels_on_path(
        state: RSAGNModelEnvState, path: chex.Array, slot_index: chex.Array, req_slots: int, params: RSAGNModelEnvParams
) -> chex.Array:
    """Get the minimum value of the SNR on newly assigned channels.
    N.B. this requires the link_snr_array to have already been calculated and present in state."""
    snr_value_all_channels = get_snr_for_path(path, state.link_snr_array, params)
    min_snr_value_sub_channels = jnp.min(
        jnp.concatenate([
            snr_value_all_channels[slot_index].reshape((1,)),
            snr_value_all_channels[slot_index + req_slots - 1].reshape((1,))
        ], axis=0)
    )
    return min_snr_value_sub_channels

get_num_spectral_features(n_nodes)

Heuristic for number of spectral features based on graph size.

Parameters:

Name	Type	Description	Default
n_nodes	int	Number of nodes in the graph	required

Returns:

Type	Description
int	Number of spectral features to use, clamped between 3 and 15.
int	Follows log2(n_nodes) scaling as reasonable default.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_num_spectral_features(n_nodes: int) -> int:
    """Heuristic for number of spectral features based on graph size.

    Args:
        n_nodes: Number of nodes in the graph

    Returns:
        Number of spectral features to use, clamped between 3 and 15.
        Follows log2(n_nodes) scaling as reasonable default.
    """
    return jnp.minimum(jnp.maximum(3, jnp.floor(jnp.log2(n_nodes))), 15).astype(int)

get_path_from_path_index_array(path_index_array, path_link_array)

Get path from path index array. Args: path_index_array (chex.Array): Path index array path_link_array (chex.Array): Path link array

Returns:

Type	Description
Array	jnp.array: path index values replaced with binary path-link arrays

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_path_from_path_index_array(path_index_array: chex.Array, path_link_array: chex.Array) -> chex.Array:
    """Get path from path index array.
    Args:
        path_index_array (chex.Array): Path index array
        path_link_array (chex.Array): Path link array

    Returns:
        jnp.array: path index values replaced with binary path-link arrays
    """
    # TODO - support unpacking bits (if this function ends up being used)
    def get_index_from_link(link):
        return jax.vmap(lambda x: path_link_array[x], in_axes=(0,))(link)

    return jax.vmap(get_index_from_link, in_axes=(0,))(path_index_array)

get_path_index_array(params, nodes)

Indices of paths between source and destination from path array

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_path_index_array(params, nodes):
    """Indices of paths between source and destination from path array"""
    # get source and destination nodes in order (for accurate indexing of path-link array)
    source, dest = nodes.astype(LARGE_INT_DTYPE)
    i = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
    index_array = jax.lax.dynamic_slice(jnp.arange(0, params.path_link_array.shape[0], dtype=LARGE_INT_DTYPE), (i,), (params.k_paths,))
    return index_array

get_path_indices(s, d, k, N, directed=False)

Get path indices for a given source, destination and number of paths. If source > destination and the graph is directed (two fibres per link, one in each direction) then an offset is added to the index to get the path in the other direction (the offset is the total number source-dest pairs).

Parameters:

Name	Type	Description	Default
s	int	Source node index	required
d	int	Destination node index	required
k	int	Number of paths	required
N	int	Number of nodes	required
directed	bool	Whether graph is directed. Defaults to False.	False

Returns:

Type	Description
Array	jnp.array: Start index on path-link array for candidate paths

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2, 3, 4))
def get_path_indices(s: int, d: int, k: int, N: int, directed: bool = False) -> chex.Array:
    """Get path indices for a given source, destination and number of paths.
    If source > destination and the graph is directed (two fibres per link, one in each direction) then an offset is
    added to the index to get the path in the other direction (the offset is the total number source-dest pairs).

    Args:
        s (int): Source node index
        d (int): Destination node index
        k (int): Number of paths
        N (int): Number of nodes
        directed (bool, optional): Whether graph is directed. Defaults to False.

    Returns:
        jnp.array: Start index on path-link array for candidate paths
    """
    node_indices = jnp.arange(N, dtype=LARGE_INT_DTYPE)
    indices_to_s = jnp.where(node_indices < s, node_indices, jnp.array(0, dtype=LARGE_INT_DTYPE))
    indices_to_d = jnp.where(node_indices < d, node_indices, jnp.array(0, dtype=LARGE_INT_DTYPE))
    # If two fibres per link, add offset to index to get fibre in other direction if source > destination
    directed_offset = directed * (s > d) * N * (N - 1) * k / 2
    # The following equation is based on the combinations formula
    forward = ((N * s + d - jnp.sum(indices_to_s, promote_integers=False) - 2 * s - 1) * k)
    backward = ((N * d + s - jnp.sum(indices_to_d, promote_integers=False) - 2 * d - 1) * k)
    return forward * (s < d) + backward * (s > d) + directed_offset.astype(LARGE_INT_DTYPE)

get_path_slots(link_slot_array, params, nodes_sd, i, agg_func='max')

Get slots on each constitutent link of path from link_slot_array (L x S), then aggregate to get (S x 1) representation of slots on path.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	link-slot array	required
params	EnvParams	environment parameters	required
nodes_sd	Array	source-destination nodes	required
i	int	path index	required
agg_func	str	aggregation function (max or sum). If max, result will be available slots on path. If sum, result will contain information on edge features.	'max'

Returns:

Name	Type	Description
slots	Array	slots on path

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 4))
def get_path_slots(link_slot_array: chex.Array, params: EnvParams, nodes_sd: chex.Array, i: int, agg_func: str = "max") -> chex.Array:
    """Get slots on each constitutent link of path from link_slot_array (L x S),
    then aggregate to get (S x 1) representation of slots on path.

    Args:
        link_slot_array: link-slot array
        params: environment parameters
        nodes_sd: source-destination nodes
        i: path index
        agg_func: aggregation function (max or sum).
            If max, result will be available slots on path.
            If sum, result will contain information on edge features.

    Returns:
        slots: slots on path
    """
    path = get_paths(params, nodes_sd)[i]
    path = path.reshape((params.num_links, 1))
    # Get links and collapse to single dimension
    num_slots = params.link_resources if agg_func == "max" else math.ceil(params.link_resources/params.aggregate_slots)
    slots = jnp.where(path, link_slot_array, jnp.zeros(num_slots, dtype=LARGE_FLOAT_DTYPE))
    # Make any -1s positive then get max for each slot across links
    if agg_func == "max":
        # Use this for getting slots from link_slot_array
        slots = jnp.max(jnp.absolute(slots), axis=0)
    elif agg_func == "sum":
        # TODO - consider using an RNN (or S5) to aggregate edge features
        # Use this (or alternative) for aggregating edge features from GNN
        slots = jnp.sum(slots, axis=0, promote_integers=False)
    elif agg_func == "mean":
        # Use this for getting mean value in slot index along path
        slots = jnp.mean(slots, axis=0)
    else:
        raise ValueError("agg_func must be 'max' or 'sum' or 'mean'")
    return slots

get_paths(params, nodes)

Get k paths between source and destination

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_paths(params, nodes):
    """Get k paths between source and destination"""
    index_array = get_path_index_array(params, nodes)
    paths = jnp.take(params.path_link_array.val, index_array, axis=0)
    if params.pack_path_bits:  # Unpack the bit-packed paths
        paths = jnp.unpackbits(paths, axis=1)[:, :params.num_links]
    return paths

get_paths_obs_gn_model(state, params)

Get observation space for launch power optimization (with numerical stability).

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_paths_obs_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> chex.Array:
    # TODO - make this just show the stats from just one path at a time
    """Get observation space for launch power optimization (with numerical stability)."""
    request_array = state.request_array.reshape((-1,))
    path_stats = calculate_path_stats(state, params, request_array)
    # Remove first 3 items of path stats for each path
    path_stats = path_stats[:, 3:]
    link_length_array = jnp.sum(params.link_length_array.val, axis=1, promote_integers=False)
    lightpath_snr_array = get_lightpath_snr(state, params)
    nodes_sd, requested_datarate = read_rsa_request(request_array)
    source, dest = nodes_sd

    def calculate_gn_path_stats(k_path_index, init_val):
        # Get path index
        path_index = get_path_indices(source, dest, params.k_paths, params.num_nodes,
                                      directed=params.directed_graph) + k_path_index
        path_link_array = jnp.unpackbits(params.path_link_array.val, axis=1)[:, :params.num_links] if params.pack_path_bits \
            else params.path_link_array.val
        path = path_link_array[path_index]
        path_length = jnp.dot(path, link_length_array)
        max_path_length = jnp.max(jnp.dot(path_link_array, link_length_array))
        path_length_norm = path_length / max_path_length
        max_path_length_hops = jnp.max(jnp.sum(path_link_array, axis=1, promote_integers=False))
        path_length_hops_norm = jnp.sum(path, promote_integers=False).astype(LARGE_FLOAT_DTYPE) / max_path_length_hops
        # Connections on path
        num_connections = jnp.where(path == 1, jnp.where(state.channel_power_array > 0, one, zero).sum(axis=1), zero).sum()
        num_connections_norm = num_connections / jnp.array(params.link_resources, dtype=LARGE_FLOAT_DTYPE)
        # Mean power of connections on path
        # make path with row length equal to link_resource (+1 to avoid zero division)
        mean_power_norm = (jnp.where(path == one, state.channel_power_array.sum(axis=1), zero).sum() /
                           (jnp.where(num_connections > zero, num_connections, one) * params.max_power))
        # Mean SNR of connections on the path links
        max_snr = jnp.array(50, dtype=LARGE_FLOAT_DTYPE)  # Nominal value for max GSNR in dB
        mean_snr_norm = (jnp.where(path == one, lightpath_snr_array.sum(axis=1), zero).sum(promote_integers=False) /
                         (jnp.where(num_connections > zero, num_connections, one) * max_snr))
        return jax.lax.dynamic_update_slice(
            init_val,
            jnp.array([[
                path_length,
                path_length_hops_norm,
                num_connections_norm,
                mean_power_norm,
                mean_snr_norm
            ]]),
            (k_path_index, 0),
        )

    gn_path_stats = jnp.zeros((params.k_paths, 5), dtype=LARGE_FLOAT_DTYPE)
    gn_path_stats = jax.lax.fori_loop(
        0, params.k_paths, calculate_gn_path_stats, gn_path_stats
    )
    all_stats = jnp.concatenate([path_stats, gn_path_stats], axis=1)
    return jnp.concatenate(
        (
            jnp.array([source]),
            requested_datarate / 100.,
            jnp.array([dest]),
            jnp.reshape(state.holding_time, (-1,)),
            jnp.reshape(all_stats, (-1,)),
        ),
        axis=0,
    )

get_paths_se(params, nodes)

Get max. spectral efficiency of modulation format on k paths between source and destination

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def get_paths_se(params, nodes):
    """Get max. spectral efficiency of modulation format on k paths between source and destination"""
    # get source and destination nodes in order (for accurate indexing of path-link array)
    index_array = get_path_index_array(params, nodes)
    return jnp.take(params.path_se_array.val, index_array, axis=0)

get_required_snr_se_kurtosis_array(modulation_format_index_array, col_index, params)

Convert modulation format index to required SNR or spectral efficiency. Modulation format index array contains the index of the modulation format used by the channel. The modulation index references a row in the modulations array, which contains SNR and SE values.

Parameters:

Name	Type	Description	Default
modulation_format_index_array	Array	Modulation format index array	required
col_index	int	Column index for required SNR or spectral efficiency	required
params	RSAGNModelEnvParams	Environment parameters	required

Returns:

Type	Description
Array	jnp.array: Required SNR for each channel (min. SNR for empty channel (mod. index 0))

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 2,))
def get_required_snr_se_kurtosis_array(modulation_format_index_array: chex.Array, col_index: int, params: RSAGNModelEnvParams) -> chex.Array:
    """Convert modulation format index to required SNR or spectral efficiency.
    Modulation format index array contains the index of the modulation format used by the channel.
    The modulation index references a row in the modulations array, which contains SNR and SE values.

    Args:
        modulation_format_index_array (chex.Array): Modulation format index array
        col_index (int): Column index for required SNR or spectral efficiency
        params (RSAGNModelEnvParams): Environment parameters

    Returns:
        jnp.array: Required SNR for each channel (min. SNR for empty channel (mod. index 0))
    """
    return jax.vmap(get_required_snr_se_kurtosis_on_link, in_axes=(0, None, None))(modulation_format_index_array, col_index, params)

get_snr_link_array(state, params)

Get SNR per link Args: state (EnvState): Environment state params (EnvParams): Environment parameters Returns: jnp.array: SNR per link

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def get_snr_link_array(state: EnvState, params: EnvParams) -> chex.Array:
    """Get SNR per link
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: SNR per link
    """

    def get_link_snr(link_index, state, params):
        # Get channel power, channel centre, bandwidth, and noise figure
        link_lengths = params.link_length_array[link_index, :]
        num_spans = jnp.ceil(jnp.sum(link_lengths)*1e3 / params.max_span_length).astype(MED_INT_DTYPE)
        if params.mod_format_correction:
            mod_format_link = state.modulation_format_index_array[link_index, :]
            kurtosis_link = get_required_snr_se_kurtosis_on_link(mod_format_link, 4, params)
            se_link = get_required_snr_se_kurtosis_on_link(mod_format_link, 1, params)
        else:
            kurtosis_link = jnp.zeros(params.link_resources)
            se_link = jnp.ones(params.link_resources)
        bw_link = state.channel_centre_bw_array[link_index, :]
        ch_power_link = state.channel_power_array[link_index, :]
        required_slots_link = get_required_slots_on_link(bw_link, se_link, params)
        ch_centres_link = get_centre_freq_on_link(jnp.arange(params.link_resources), required_slots_link, params)

        # Calculate SNR
        P = dict(
            num_channels=params.link_resources,
            num_spans=num_spans,
            max_spans=params.max_spans,
            ref_lambda=params.ref_lambda,
            length=link_lengths,
            attenuation_i=jnp.array(params.attenuation),
            attenuation_bar_i=jnp.array(params.attenuation_bar),
            nonlinear_coeff=jnp.array(params.nonlinear_coeff),
            raman_gain_slope_i=jnp.array(params.raman_gain_slope),
            dispersion_coeff=jnp.array(params.dispersion_coeff),
            dispersion_slope=jnp.array(params.dispersion_slope),
            coherent=params.coherent,
            num_roadms=params.num_roadms,
            roadm_loss=params.roadm_loss,
            amplifier_noise_figure=params.amplifier_noise_figure.val,
            transceiver_snr=params.transceiver_snr.val,
            mod_format_correction=params.mod_format_correction,
            ch_power_w_i=ch_power_link,
            ch_centre_i=ch_centres_link*1e9,
            ch_bandwidth_i=bw_link*1e9,
            excess_kurtosis_i=kurtosis_link,
            uniform_spans=params.uniform_spans,
        )
        snr = isrs_gn_model.get_snr(**P)[0]

        return snr

    link_snr_array = jax.vmap(get_link_snr, in_axes=(0, None, None))(jnp.arange(params.num_links), state, params)
    link_snr_array = jnp.nan_to_num(link_snr_array, nan=1e-5)
    return link_snr_array

get_spectral_features(laplacian, num_features)

Compute spectral node features from symmetric normalized graph Laplacian.

Parameters:

Name	Type	Description	Default
adj		Adjacency matrix of the graph	required
num_features	int	Number of eigenvector features to extract	required

Returns:

Type	Description
ndarray	Array of shape (n_nodes, num_features) containing eigenvectors corresponding
ndarray	to the smallest non-zero eigenvalues of the graph Laplacian.

Notes

• Skips trivial eigenvectors (those with near-zero eigenvalues)
• Eigenvectors are ordered by ascending eigenvalue magnitude
• Runtime is O(n^3) - use only for small/medium graphs
• Eigenvector signs are arbitrary (may vary between runs)

Source code in xlron/environments/env_funcs.py

def get_spectral_features(laplacian: jnp.array, num_features: int) -> jnp.ndarray:
    """Compute spectral node features from symmetric normalized graph Laplacian.

    Args:
        adj: Adjacency matrix of the graph
        num_features: Number of eigenvector features to extract

    Returns:
        Array of shape (n_nodes, num_features) containing eigenvectors corresponding
        to the smallest non-zero eigenvalues of the graph Laplacian.

    Notes:
        - Skips trivial eigenvectors (those with near-zero eigenvalues)
        - Eigenvectors are ordered by ascending eigenvalue magnitude
        - Runtime is O(n^3) - use only for small/medium graphs
        - Eigenvector signs are arbitrary (may vary between runs)
    """
    n_nodes = laplacian.shape[0]
    eigenvalues, eigenvectors = jnp.linalg.eigh(laplacian)
    return eigenvectors[:, :num_features].astype(LARGE_FLOAT_DTYPE)

implement_action_rmsa_gn_model(state, action, params)

Implement action for RSA GN model. Update following arrays: - link_slot_array - link_slot_departure_array - link_snr_array - modulation_format_index_array - channel_power_array - active_path_array Args: state (EnvState): Environment state action (chex.Array): Action tuple (first is path action, second is launch_power) params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rmsa_gn_model(
        state: RSAGNModelEnvState, action: chex.Array, params: RSAGNModelEnvParams
) -> EnvState:
    """Implement action for RSA GN model. Update following arrays:
    - link_slot_array
    - link_slot_departure_array
    - link_snr_array
    - modulation_format_index_array
    - channel_power_array
    - active_path_array
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action tuple (first is path action, second is launch_power)
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_action, power_action = action
    path_action = path_action.astype(MED_INT_DTYPE)
    k_path_index, initial_slot_index = process_path_action(state, params, path_action)
    lightpath_index = get_lightpath_index(params, nodes_sd, k_path_index)
    path = get_paths(params, nodes_sd)[k_path_index]
    launch_power = get_launch_power(state, path_action, power_action, params)
    # TODO(GN MODEL) - get mod. format based on maximum reach
    mod_format_index = jax.lax.dynamic_slice(
        state.mod_format_mask, (path_action,), (1,)
    ).astype(MED_INT_DTYPE)[0]
    se = params.modulations_array.val[mod_format_index][1]
    num_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)
    # Update link_slot_array and link_slot_departure_array, then other arrays
    state = implement_path_action(state, path, initial_slot_index, num_slots)
    state = state.replace(
        path_index_array=vmap_set_path_links(state.path_index_array, path, initial_slot_index, num_slots-params.guardband, lightpath_index),
        channel_power_array=vmap_set_path_links(state.channel_power_array, path, initial_slot_index, num_slots-params.guardband, launch_power),
        modulation_format_index_array=vmap_set_path_links(state.modulation_format_index_array, path, initial_slot_index, num_slots-params.guardband, mod_format_index),
        channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, initial_slot_index, num_slots-params.guardband, params.slot_size),
    )
    # Update link_snr_array
    state = state.replace(link_snr_array=get_snr_link_array(state, params))
    # jax.debug.print("launch_power {}", launch_power, ordered=True)
    # jax.debug.print("mod_format_index {}", mod_format_index, ordered=True)
    # jax.debug.print("initial_slot_index {}", initial_slot_index, ordered=True)
    # jax.debug.print("state.mod_format_mask {}", state.mod_format_mask, ordered=True)
    # jax.debug.print("path_snr {}", get_snr_for_path(path, state.link_snr_array, params), ordered=True)
    # jax.debug.print("required_snr {}", params.modulations_array.val[mod_format_index][2] + params.snr_margin, ordered=True)
    return state

implement_action_rsa(state, action, params)

Implement action to assign slots on links.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required
action	Array	action to implement	required
params	EnvParams	environment parameters	required

Returns:

Name	Type	Description
state	EnvState	updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rsa(
        state: EnvState,
        action: chex.Array,
        params: EnvParams,
) -> EnvState:
    """Implement action to assign slots on links.

    Args:
        state: current state
        action: action to implement
        params: environment parameters

    Returns:
        state: updated state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_index, initial_slot_index = process_path_action(state, params, action)
    path = get_paths(params, nodes_sd)[path_index]
    if params.__class__.__name__ == "RWALightpathReuseEnvParams":
        state = state.replace(
            link_capacity_array=vmap_update_path_links(
                state.link_capacity_array, path, initial_slot_index, 1, requested_datarate
            )
        )
        # TODO (Dynamic-RWALR) - to support diverse requested_datarates for RWA-LR, need to update masking
        # TODO (Dynamic-RWALR) - In order to enable dynamic RWA with lightpath reuse (as opposed to just incremental loading),
        #  need to keep track of active requests OR just randomly remove connections
        #  (could do this by using the link_slot_departure array in a novel way... i.e. don't fill it with departure time but current bw)
        capacity_mask = jnp.where(state.link_capacity_array <= 0., -1., 0.)
        over_capacity_mask = jnp.where(state.link_capacity_array < 0., -1., 0.)
        total_mask = capacity_mask + over_capacity_mask
        state = state.replace(
            link_slot_array=total_mask,
            link_slot_departure_array=vmap_update_path_links(state.link_slot_departure_array, path,
                                                                       initial_slot_index, 1,
                                                                       state.current_time + state.holding_time)
        )
    else:
        se = get_paths_se(params, nodes_sd)[path_index] if params.consider_modulation_format else one
        num_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)
        state = implement_path_action(state, path, initial_slot_index, num_slots)
    return state

implement_action_rsa_gn_model(state, action, params)

Implement action for RSA GN model. Update following arrays: - link_slot_array - link_slot_departure_array - link_snr_array - modulation_format_index_array - channel_power_array - active_path_array Args: state (EnvState): Environment state action (chex.Array): Action tuple (first is path action, second is launch_power) params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rsa_gn_model(
        state: RSAGNModelEnvState, action: chex.Array, params: RSAGNModelEnvParams
) -> EnvState:
    """Implement action for RSA GN model. Update following arrays:
    - link_slot_array
    - link_slot_departure_array
    - link_snr_array
    - modulation_format_index_array
    - channel_power_array
    - active_path_array
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action tuple (first is path action, second is launch_power)
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_action, power_action = action
    path_action = path_action.astype(MED_INT_DTYPE)
    k_path_index, initial_slot_index = process_path_action(state, params, path_action)
    lightpath_index = get_lightpath_index(params, nodes_sd, k_path_index)
    path = get_paths(params, nodes_sd)[k_path_index]
    launch_power = get_launch_power(state, path_action, power_action, params)
    num_slots = required_slots(requested_datarate, 1, params.slot_size, guardband=params.guardband)
    # Update link_slot_array and link_slot_departure_array, then other arrays
    state = implement_path_action(state, path, initial_slot_index, num_slots)
    state = state.replace(
        path_index_array=vmap_set_path_links(state.path_index_array, path, initial_slot_index, num_slots-params.guardband, lightpath_index),
        channel_power_array=vmap_set_path_links(state.channel_power_array, path, initial_slot_index, num_slots-params.guardband, launch_power),
        # TODO - update this to use separate arrays to track channel centres and bandwidths and update with bandwidth (that may or may not equal slot size)
        channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, initial_slot_index, num_slots-params.guardband, params.slot_size),
    )
    if params.monitor_active_lightpaths:
        state = state.replace(
            active_lightpaths_array=update_active_lightpaths_array(state, lightpath_index, initial_slot_index, num_slots-params.guardband),
            active_lightpaths_array_departure=update_active_lightpaths_array_departure(state, -state.current_time-state.holding_time),
        )
        # No need to check SNR until end of episode
        return state
    # Update link_snr_array
    state = state.replace(link_snr_array=get_snr_link_array(state, params))
    return state

implement_action_rwalr(state, action, params)

For use in RWALightpathReuseEnv. Update link_slot_array and link_slot_departure_array to reflect new lightpath assignment. Update link_capacity_array with new capacity if lightpath is available. Undo link_capacity_update if over capacity.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
action	Array	Action array	required
params	EnvParams	Environment parameters	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2,))
def implement_action_rwalr(state: EnvState, action: chex.Array, params: EnvParams) -> EnvState:
    """For use in RWALightpathReuseEnv.
    Update link_slot_array and link_slot_departure_array to reflect new lightpath assignment.
    Update link_capacity_array with new capacity if lightpath is available.
    Undo link_capacity_update if over capacity.

    Args:
        state: Environment state
        action: Action array
        params: Environment parameters

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(state.request_array)
    path_index, initial_slot_index = process_path_action(state, params, action)
    path = get_paths(params, nodes_sd)[path_index]
    lightpath_available_check, lightpath_existing_check, curr_lightpath_capacity, lightpath_index = (
        check_lightpath_available_and_existing(state, params, action)
    )
    # Get path capacity - request
    lightpath_capacity = jax.lax.cond(
        lightpath_existing_check,
        lambda x: curr_lightpath_capacity - requested_datarate,  # Subtract requested_datarate from current lightpath
        lambda x: jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.path_capacity_array, x, 1)) - requested_datarate,  # Get initial capacity of lightpath - request
        lightpath_index
    )
    # Update link_capacity_array with new capacity if lightpath is available
    state = jax.lax.cond(
        lightpath_available_check,
        lambda x: x.replace(
            link_capacity_array=vmap_set_path_links(
                state.link_capacity_array, path, initial_slot_index, 1, lightpath_capacity
            ),
            path_index_array=vmap_set_path_links(
                state.path_index_array, path, initial_slot_index, 1, lightpath_index
            ),
        ),
        lambda x: x,
        state
    )
    capacity_mask = jnp.where(state.link_capacity_array <= 0., -1., 0.)
    over_capacity_mask = jnp.where(state.link_capacity_array < 0., -1., 0.)
    # Undo link_capacity_update if over capacity
    # N.B. this will fail if requested capacity is greater than total original capacity of lightpath
    lightpath_capacity_before_action = jax.lax.cond(
        lightpath_existing_check,
        lambda x: curr_lightpath_capacity,  # Subtract requested_datarate from current lightpath
        lambda x: 1e6,  # Empty slots have high capacity (1e6)
        # Get initial capacity of lightpath - request
        None,
    )
    state = state.replace(
        link_capacity_array=jnp.where(over_capacity_mask == -1, lightpath_capacity_before_action, state.link_capacity_array)
    )
    # Total mask will be 0 if space still available, -1 if capacity is zero or -2 if over capacity
    total_mask = capacity_mask + over_capacity_mask
    # Update link_slot_array and link_slot_departure_array
    state = state.replace(
        link_slot_array=total_mask,
        link_slot_departure_array=vmap_update_path_links(state.link_slot_departure_array, path,
                                                                 initial_slot_index, 1,
                                                                 state.current_time + state.holding_time)
    )
    return state

implement_node_action(state, s_node, d_node, s_request, d_request, n=2)

Update node capacity, node resource and node departure arrays

Parameters:

Name	Type	Description	Default
state	State	current state	required
s_node	int	source node	required
d_node	int	destination node	required
s_request	int	source node request	required
d_request	int	destination node request	required
n	int	number of nodes to implement. Defaults to 2.	2

Returns:

Name	Type	Description
State	EnvState	updated state

Source code in xlron/environments/env_funcs.py

def implement_node_action(state: EnvState, s_node: chex.Array, d_node: chex.Array, s_request: chex.Array, d_request: chex.Array, n=2) -> EnvState:
    """Update node capacity, node resource and node departure arrays

    Args:
        state (State): current state
        s_node (int): source node
        d_node (int): destination node
        s_request (int): source node request
        d_request (int): destination node request
        n (int, optional): number of nodes to implement. Defaults to 2.

    Returns:
        State: updated state
    """
    node_indices = jnp.arange(state.node_capacity_array.shape[0])

    curr_selected_nodes = jnp.zeros(state.node_capacity_array.shape[0])
    # d_request -ve so that selected node is +ve (so that argmin works correctly for node resource array update)
    # curr_selected_nodes is N x 1 array, with requested node resources at index of selected node
    curr_selected_nodes = update_node_array(node_indices, curr_selected_nodes, d_node, -d_request)
    curr_selected_nodes = jax.lax.cond(n == 2, lambda x: update_node_array(*x), lambda x: x[1], (node_indices, curr_selected_nodes, s_node, -s_request))

    node_capacity_array = state.node_capacity_array - curr_selected_nodes

    node_resource_array = vmap_update_node_resources(state.node_resource_array, curr_selected_nodes)

    node_departure_array = vmap_update_node_departure(state.node_departure_array, curr_selected_nodes, -state.current_time-state.holding_time)

    state = state.replace(
        node_capacity_array=node_capacity_array,
        node_resource_array=node_resource_array,
        node_departure_array=node_departure_array
    )

    return state

implement_path_action(state, path, initial_slot_index, num_slots)

Update link-slot and link-slot departure arrays. Times are set to negative until turned positive by finalisation (after checks).

Parameters:

Name	Type	Description	Default
state	State	current state	required
path	int	path to implement	required
initial_slot_index	int	initial slot index	required
num_slots	int	number of slots to implement	required

Source code in xlron/environments/env_funcs.py

def implement_path_action(state: EnvState, path: chex.Array, initial_slot_index: chex.Array, num_slots: chex.Array) -> EnvState:
    """Update link-slot and link-slot departure arrays.
    Times are set to negative until turned positive by finalisation (after checks).

    Args:
        state (State): current state
        path (int): path to implement
        initial_slot_index (int): initial slot index
        num_slots (int): number of slots to implement
    """
    state = state.replace(
        link_slot_array=vmap_update_path_links(state.link_slot_array, path, initial_slot_index, num_slots, one),
        link_slot_departure_array=vmap_update_path_links(state.link_slot_departure_array, path, initial_slot_index, num_slots, state.current_time+state.holding_time)
    )
    return state

implement_vone_action(state, action, total_actions, remaining_actions, params)

Implement action to assign nodes (1, 2, or 0 nodes assigned per action) and assign slots and links for lightpath.

Parameters:

Name	Type	Description	Default
state	EnvState	current state	required
action	Array	action to implement (node, node, path_slot_action)	required
total_actions	Scalar	total number of actions to implement for current request	required
remaining_actions	Scalar	remaining actions to implement	required
k		number of paths to consider	required
N		number of nodes to assign	required

Returns:

Name	Type	Description
state		updated state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(4,))
def implement_vone_action(
        state: EnvState,
        action: chex.Array,
        total_actions: chex.Scalar,
        remaining_actions: chex.Scalar,
        params: EnvParams,
):
    """Implement action to assign nodes (1, 2, or 0 nodes assigned per action) and assign slots and links for lightpath.

    Args:
        state: current state
        action: action to implement (node, node, path_slot_action)
        total_actions: total number of actions to implement for current request
        remaining_actions: remaining actions to implement
        k: number of paths to consider
        N: number of nodes to assign

    Returns:
        state: updated state
    """
    request = jax.lax.dynamic_slice(state.request_array[0], ((remaining_actions-1)*2, ), (3, ))
    node_request_s = jax.lax.dynamic_slice(request, (2, ), (1, ))
    requested_datarate = jax.lax.dynamic_slice(request, (1,), (1,))
    node_request_d = jax.lax.dynamic_slice(request, (0, ), (1, ))
    nodes = action[::2]
    path_index, initial_slot_index = process_path_action(state, params, action[1])
    path = get_paths(params, nodes)[path_index]
    se = get_paths_se(params, nodes)[path_index] if params.consider_modulation_format else jnp.array([1])
    num_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)

    # jax.debug.print("state.request_array {}", state.request_array, ordered=True)
    # jax.debug.print("path {}", path, ordered=True)
    # jax.debug.print("slots {}", jnp.max(jnp.where(path.reshape(-1,1) == 1, state.link_slot_array, jnp.zeros(params.num_links).reshape(-1,1)), axis=0), ordered=True)
    # jax.debug.print("path_index {}", path_index, ordered=True)
    # jax.debug.print("initial_slot_index {}", initial_slot_index, ordered=True)
    # jax.debug.print("requested_datarate {}", requested_datarate, ordered=True)
    # jax.debug.print("request {}", request, ordered=True)
    # jax.debug.print("se {}", se, ordered=True)
    # jax.debug.print("num_slots {}", num_slots, ordered=True)

    n_nodes = jax.lax.cond(
        total_actions == remaining_actions,
        lambda x: 2, lambda x: 1,
        (total_actions, remaining_actions))
    path_action_only_check = path_action_only(state.request_array[1], state.action_counter, remaining_actions)

    state = jax.lax.cond(
        path_action_only_check,
        lambda x: x[0],
        lambda x: implement_node_action(x[0], x[1], x[2], x[3], x[4], n=x[5]),
        (state, nodes[0], nodes[1], node_request_s, node_request_d, n_nodes)
    )

    state = implement_path_action(state, path, initial_slot_index, num_slots)

    return state

init_action_counter()

Initialize action counter. First index is num unique nodes, second index is total steps, final is remaining steps until completion of request. Only used in VONE environments.

Source code in xlron/environments/env_funcs.py

def init_action_counter():
    """Initialize action counter.
    First index is num unique nodes, second index is total steps, final is remaining steps until completion of request.
    Only used in VONE environments.
    """
    return jnp.zeros(3, dtype=MED_INT_DTYPE)

init_action_history(params)

Initialize action history

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_action_history(params: EnvParams):
    """Initialize action history"""
    return jnp.full(params.max_edges*2+1, -1, dtype=LARGE_FLOAT_DTYPE)

init_active_lightpaths_array(params)

Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array. M is MIN(max_requests, num_links * link_resources / min_slots). min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

Parameters:

Name	Type	Description	Default
params	RSAGNModelEnvParams	Environment parameters	required

Returns: jnp.array: Active path array (default value -1, empty path)

Source code in xlron/environments/env_funcs.py

def init_active_lightpaths_array(params: RSAGNModelEnvParams):
    """Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array.
    M is MIN(max_requests, num_links * link_resources / min_slots).
    min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

    Args:
        params (RSAGNModelEnvParams): Environment parameters
    Returns:
        jnp.array: Active path array (default value -1, empty path)
    """
    total_slots = params.num_links * params.link_resources  # total slots on networks
    min_slots = jnp.max(params.values_bw.val) / params.slot_size  # minimum number of slots required for lightpath
    return jnp.full((int(total_slots / min_slots), 3), -1, dtype=LARGE_INT_DTYPE)

init_active_lightpaths_array_departure(params)

Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array. M is MIN(max_requests, num_links * link_resources / min_slots). min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

Parameters:

Name	Type	Description	Default
params	RSAGNModelEnvParams	Environment parameters	required

Returns: jnp.array: Active path array (default value -1, empty path)

Source code in xlron/environments/env_funcs.py

def init_active_lightpaths_array_departure(params: RSAGNModelEnvParams):
    """Initialise active lightpath array. Stores path indices of all active paths on the network in a 1 x M array.
    M is MIN(max_requests, num_links * link_resources / min_slots).
    min_slots is the minimum number of slots required for a lightpath i.e. max(values_bw)/ slot_size.

    Args:
        params (RSAGNModelEnvParams): Environment parameters
    Returns:
        jnp.array: Active path array (default value -1, empty path)
    """
    total_slots = params.num_links * params.link_resources  # total slots on networks
    min_slots = jnp.max(params.values_bw.val) / params.slot_size  # minimum number of slots required for lightpath
    return jnp.full((int(total_slots / min_slots), 3), 0., dtype=SMALL_FLOAT_DTYPE)

init_active_path_array(params)

Initialise active path array. Stores details of full path utilised by lightpath on each frequency slot. Args: params (EnvParams): Environment parameters Returns: jnp.array: Active path array

Source code in xlron/environments/env_funcs.py

def init_active_path_array(params: EnvParams):
    """Initialise active path array. Stores details of full path utilised by lightpath on each frequency slot.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Active path array
    """
    return jnp.full((params.num_links, params.link_resources, params.num_links), -1, dtype=MED_INT_DTYPE)

init_channel_centre_bw_array(params)

Initialise channel centre array. Args: params (EnvParams): Environment parameters Returns: jnp.array: Channel centre array

Source code in xlron/environments/env_funcs.py

def init_channel_centre_bw_array(params: EnvParams):
    """Initialise channel centre array.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Channel centre array
    """
    return jnp.full((params.num_links, params.link_resources), 0., dtype=LARGE_FLOAT_DTYPE)

init_channel_power_array(params)

Initialise channel power array.

Parameters:

Name	Type	Description	Default
params	EnvParams	Environment parameters	required

Returns: jnp.array: Channel power array

Source code in xlron/environments/env_funcs.py

def init_channel_power_array(params: EnvParams):
    """Initialise channel power array.

    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Channel power array
    """
    return jnp.full((params.num_links, params.link_resources), 0., dtype=LARGE_FLOAT_DTYPE)

init_graph_tuple(state, params, adj, exclude_source_dest=False)

Initialise graph tuple for use with Jraph GNNs. Args: state (EnvState): Environment state params (EnvParams): Environment parameters adj (jnp.array): Adjacency matrix of the graph Returns: jraph.GraphsTuple: Graph tuple

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 3))
def init_graph_tuple(state: EnvState, params: EnvParams, adj: jnp.array, exclude_source_dest: bool=False) -> jraph.GraphsTuple:
    """Initialise graph tuple for use with Jraph GNNs.
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        adj (jnp.array): Adjacency matrix of the graph
    Returns:
        jraph.GraphsTuple: Graph tuple
    """
    senders = params.edges.val.T[0]
    receivers = params.edges.val.T[1]

    # Get source and dest from request array
    source_dest, datarate = read_rsa_request(state.request_array)
    # Global feature is normalised data rate of current request
    globals = jnp.array([datarate / jnp.max(params.values_bw.val)], dtype=LARGE_FLOAT_DTYPE)

    if exclude_source_dest:
        source_dest_features = jnp.zeros((params.num_nodes, 2), dtype=LARGE_FLOAT_DTYPE)
    else:
        source, dest = source_dest[0], source_dest[2]
        # One-hot encode source and destination (2 additional features)
        source_dest_features = jnp.zeros((params.num_nodes, 2), dtype=LARGE_FLOAT_DTYPE)
        source_dest_features = source_dest_features.at[source.astype(MED_INT_DTYPE), 0].set(1)
        source_dest_features = source_dest_features.at[dest.astype(MED_INT_DTYPE), 1].set(-1)

    spectral_features = get_spectral_features(adj, num_features=3)

    # For dynamic traffic, edge_features are normalised remaining holding time instead of link_slot_array
    holding_time_edge_features = state.link_slot_departure_array / params.mean_service_holding_time

    if params.__class__.__name__ in ["RSAGNModelEnvParams", "RMSAGNModelEnvParams"]:
        # Normalize by max parameters (converted to linear units)
        max_power = isrs_gn_model.from_dbm(params.max_power)
        normalized_power = jnp.round(state.channel_power_array / max_power, 3)
        max_snr = isrs_gn_model.from_db(params.max_snr)
        normalized_snr = jnp.round(state.link_snr_array / max_snr, 3)
        edge_features = jnp.stack([normalized_snr, normalized_power], axis=-1)
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)
    elif params.__class__.__name__ == "VONEEnvParams":
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        node_features = getattr(state, "node_capacity_array", jnp.zeros(params.num_nodes, dtype=LARGE_FLOAT_DTYPE))
        node_features = node_features.reshape(-1, 1)
        node_features = jnp.concatenate([node_features, spectral_features, source_dest_features], axis=-1)
    else:
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        # [n_edges] or [n_edges, ...]
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)

    if params.disable_node_features:
        node_features = jnp.zeros((1,), dtype=LARGE_FLOAT_DTYPE)

    # Handle undirected graphs (duplicate edges after normalization)
    if not params.directed_graph:
        senders_ = jnp.concatenate([senders, receivers])
        receivers = jnp.concatenate([receivers, senders])
        senders = senders_
        edge_features = jnp.repeat(edge_features, 2, axis=0)

    return jraph.GraphsTuple(
        nodes=node_features,
        edges=edge_features,
        senders=senders,
        receivers=receivers,
        n_node=jnp.reshape(params.num_nodes, (1,)),
        n_edge=jnp.reshape(len(senders), (1,)),
        globals=globals,
    )

init_link_capacity_array(params)

Initialise link capacity array. Represents available data rate for lightpath on each link. Default is high value (1e6) for unoccupied slots. Once lightpath established, capacity is determined by corresponding entry in path capacity array.

Source code in xlron/environments/env_funcs.py

def init_link_capacity_array(params):
    """Initialise link capacity array. Represents available data rate for lightpath on each link.
    Default is high value (1e6) for unoccupied slots. Once lightpath established, capacity is determined by
    corresponding entry in path capacity array."""
    return jnp.full((params.num_links, params.link_resources), 1e6)

init_link_length_array(graph)

Initialise link length array. Args: graph (nx.Graph): NetworkX graph Returns:

Source code in xlron/environments/env_funcs.py

def init_link_length_array(graph: nx.Graph) -> chex.Array:
    """Initialise link length array.
    Args:
        graph (nx.Graph): NetworkX graph
    Returns:

    """
    link_lengths = []
    for edge in sorted(graph.edges):
        link_lengths.append(graph.edges[edge]["weight"])
    return jnp.array(link_lengths, dtype=MED_INT_DTYPE)

init_link_length_array_gn_model(graph, max_span_length, max_spans)

Initialise link length array for environements that use GN model of physical layer. We assume each link has spans of equal length.

Parameters:

Name	Type	Description	Default
graph	Graph	NetworkX graph	required

Returns: jnp.array: Link length array (L x max_spans)

Source code in xlron/environments/env_funcs.py

def init_link_length_array_gn_model(graph: nx.Graph, max_span_length: int,  max_spans: int) -> chex.Array:
    """Initialise link length array for environements that use GN model of physical layer.
    We assume each link has spans of equal length.

    Args:
        graph (nx.Graph): NetworkX graph
    Returns:
        jnp.array: Link length array (L x max_spans)
    """
    link_lengths = []
    directed = graph.is_directed()
    graph = graph.to_undirected()
    edges = sorted(graph.edges)
    for edge in edges:
        link_lengths.append(graph.edges[edge]["weight"])
    if directed:
        for edge in edges:
            link_lengths.append(graph.edges[edge]["weight"])
    span_length_array = []
    for length in link_lengths:
        num_spans = math.ceil(length / max_span_length)
        avg_span_length = length / num_spans
        span_lengths = [avg_span_length] * num_spans
        span_lengths.extend([0] * (max_spans - num_spans))
        span_length_array.append(span_lengths)
    return jnp.array(span_length_array, dtype=MED_INT_DTYPE)

init_link_slot_array(params)

Initialize empty (all zeroes) link-slot array. 0 means slot is free, -1 means occupied. Args: params (EnvParams): Environment parameters Returns: jnp.array: Link slot array (E x S) where E is number of edges and S is number of slots

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_link_slot_array(params: EnvParams):
    """Initialize empty (all zeroes) link-slot array. 0 means slot is free, -1 means occupied.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Link slot array (E x S) where E is number of edges and S is number of slots"""
    return jnp.zeros((params.num_links, params.link_resources), dtype=LARGE_FLOAT_DTYPE)

init_link_slot_mask(params, agg=1)

Initialize link mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0, 1))
def init_link_slot_mask(params: EnvParams, agg: int = 1):
    """Initialize link mask"""
    return jnp.ones(params.k_paths*math.ceil(params.link_resources / agg), dtype=LARGE_FLOAT_DTYPE)

init_link_snr_array(params)

Initialise signal-to-noise ratio (SNR) array. Args: params (EnvParams): Environment parameters Returns: jnp.array: SNR array

Source code in xlron/environments/env_funcs.py

def init_link_snr_array(params: EnvParams):
    """Initialise signal-to-noise ratio (SNR) array.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: SNR array
    """
    # The SNR is kept in linear units to allow summation of 1/SNR across links
    return jnp.full((params.num_links, params.link_resources), -1e5, dtype=LARGE_FLOAT_DTYPE)

init_mod_format_mask(params)

Initialize link mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_mod_format_mask(params: EnvParams):
    """Initialize link mask"""
    return jnp.full((params.k_paths*params.link_resources,), -1.0, dtype=LARGE_FLOAT_DTYPE)

init_modulation_format_index_array(params)

Initialise modulation format index array. Args: params (EnvParams): Environment parameters Returns: jnp.array: Modulation format index array

Source code in xlron/environments/env_funcs.py

def init_modulation_format_index_array(params: EnvParams):
    """Initialise modulation format index array.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Modulation format index array
    """
    return jnp.full((params.num_links, params.link_resources), -1, dtype=MED_INT_DTYPE)  # -1 so that highest order is assumed (closest to Gaussian)

init_modulations_array(modulations_filepath=None)

Initialise array of maximum spectral efficiency for modulation format on path.

Parameters:

Name	Type	Description	Default
modulations_filepath	str	Path to CSV file containing modulation formats. Defaults to None.	None

Returns: jnp.array: Array of maximum spectral efficiency for modulation format on path. First two columns are maximum path length and spectral efficiency.

Source code in xlron/environments/env_funcs.py

def init_modulations_array(modulations_filepath: str = None):
    """Initialise array of maximum spectral efficiency for modulation format on path.

    Args:
        modulations_filepath (str, optional): Path to CSV file containing modulation formats. Defaults to None.
    Returns:
        jnp.array: Array of maximum spectral efficiency for modulation format on path.
        First two columns are maximum path length and spectral efficiency.
    """
    f = pathlib.Path(modulations_filepath) if modulations_filepath else (
            pathlib.Path(__file__).parents[1].absolute() / "data" / "modulations" / "modulations.csv")
    modulations = np.genfromtxt(f, delimiter=',')
    # Drop empty first row (headers) and column (name)
    modulations = modulations[1:, 1:]
    return jnp.array(modulations, dtype=LARGE_FLOAT_DTYPE)

init_node_capacity_array(params)

Initialize node array with uniform resources. Args: params (EnvParams): Environment parameters Returns: jnp.array: Node capacity array (N x 1) where N is number of nodes

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_node_capacity_array(params: EnvParams):
    """Initialize node array with uniform resources.
    Args:
        params (EnvParams): Environment parameters
    Returns:
        jnp.array: Node capacity array (N x 1) where N is number of nodes"""
    return jnp.array([params.node_resources] * params.num_nodes, dtype=MED_INT_DTYPE)

init_node_mask(params)

Initialize node mask

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_node_mask(params: EnvParams):
    """Initialize node mask"""
    return jnp.ones(params.num_nodes, dtype=LARGE_FLOAT_DTYPE)

init_node_resource_array(params)

Array to track node resources occupied by virtual nodes

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_node_resource_array(params: EnvParams):
    """Array to track node resources occupied by virtual nodes"""
    return jnp.zeros((params.num_nodes, params.node_resources), dtype=LARGE_FLOAT_DTYPE)

init_path_capacity_array(link_length_array, path_link_array, min_request=1, scale_factor=1.0, alpha=0.0002, NF=4.5, B=10000000000000.0, R_s=100000000000.0, beta_2=-2.17e-26, gamma=0.0012, L_s=100000.0, lambda0=1.55e-06)

Calculated from Nevin paper: https://api.repository.cam.ac.uk/server/api/core/bitstreams/b80e7a9c-a86b-4b30-a6d6-05017c60b0c8/content

Parameters:

Name	Type	Description	Default
link_length_array	Array	Array of link lengths	required
path_link_array	Array	Array of links on paths	required
min_request	int	Minimum data rate request size. Defaults to 100 GBps.	1
scale_factor	float	Scale factor for link capacity. Defaults to 1.0.	1.0
alpha	float	Fibre attenuation coefficient. Defaults to 0.2e-3 /m	0.0002
NF	float	Amplifier noise figure. Defaults to 4.5 dB.	4.5
B	float	Total modulated bandwidth. Defaults to 10e12 Hz.	10000000000000.0
R_s	float	Symbol rate. Defaults to 100e9 Baud.	100000000000.0
beta_2	float	Dispersion parameter. Defaults to -21.7e-27 s^2/m.	-2.17e-26
gamma	float	Nonlinear coefficient. Defaults to 1.2e-3 /W/m.	0.0012
L_s	float	Span length. Defaults to 100e3 m.	100000.0
lambda0	float	Wavelength. Defaults to 1550e-9 m.	1.55e-06

Returns:

Type	Description
Array	chex.Array: Array of link capacities in Gbps

Source code in xlron/environments/env_funcs.py

def init_path_capacity_array(
        link_length_array: chex.Array,
        path_link_array: chex.Array,
        min_request=1,  # Minimum data rate request size
        scale_factor=1.0,  # Scale factor for link capacity
        alpha=0.2e-3,  # Fibre attenuation coefficient
        NF=4.5,  # Amplifier noise figure
        B=10e12,  # Total modulated bandwidth
        R_s=100e9,  # Symbol rate
        beta_2=-21.7e-27,  # Dispersion parameter
        gamma=1.2e-3,  # Nonlinear coefficient
        L_s=100e3,  # span length
        lambda0=1550e-9,  # Wavelength
) -> chex.Array:
    """Calculated from Nevin paper:
    https://api.repository.cam.ac.uk/server/api/core/bitstreams/b80e7a9c-a86b-4b30-a6d6-05017c60b0c8/content

    Args:
        link_length_array (chex.Array): Array of link lengths
        path_link_array (chex.Array): Array of links on paths
        min_request (int, optional): Minimum data rate request size. Defaults to 100 GBps.
        scale_factor (float, optional): Scale factor for link capacity. Defaults to 1.0.
        alpha (float, optional): Fibre attenuation coefficient. Defaults to 0.2e-3 /m
        NF (float, optional): Amplifier noise figure. Defaults to 4.5 dB.
        B (float, optional): Total modulated bandwidth. Defaults to 10e12 Hz.
        R_s (float, optional): Symbol rate. Defaults to 100e9 Baud.
        beta_2 (float, optional): Dispersion parameter. Defaults to -21.7e-27 s^2/m.
        gamma (float, optional): Nonlinear coefficient. Defaults to 1.2e-3 /W/m.
        L_s (float, optional): Span length. Defaults to 100e3 m.
        lambda0 (float, optional): Wavelength. Defaults to 1550e-9 m.

    Returns:
        chex.Array: Array of link capacities in Gbps
    """
    path_length_array = jnp.dot(path_link_array, link_length_array)
    path_capacity_array = calculate_path_capacity(
        path_length_array,
        min_request=min_request,
        scale_factor=scale_factor,
        alpha=alpha,
        NF=NF,
        B=B,
        R_s=R_s,
        beta_2=beta_2,
        gamma=gamma,
        L_s=L_s,
        lambda0=lambda0,
    )
    return path_capacity_array.astype(MED_INT_DTYPE)

init_path_index_array(params)

Initialise path index array. Represents index of lightpath occupying each slot.

Source code in xlron/environments/env_funcs.py

def init_path_index_array(params):
    """Initialise path index array. Represents index of lightpath occupying each slot."""
    return jnp.full((params.num_links, params.link_resources), -1)

init_path_length_array(path_link_array, graph)

Initialise path length array.

Parameters:

Name	Type	Description	Default
path_link_array	Array	Path-link array	required
graph	Graph	NetworkX graph	required

Returns: chex.Array: Path length array

Source code in xlron/environments/env_funcs.py

def init_path_length_array(path_link_array: chex.Array, graph: nx.Graph) -> chex.Array:
    """Initialise path length array.

    Args:
        path_link_array (chex.Array): Path-link array
        graph (nx.Graph): NetworkX graph
    Returns:
        chex.Array: Path length array
    """
    link_length_array = init_link_length_array(graph)
    path_lengths = jnp.dot(path_link_array, link_length_array)
    return path_lengths

init_path_link_array(graph, k, disjoint=False, weight='weight', directed=False, modulations_array=None, rwa_lr=False, scale_factor=1.0, path_snr=False)

Initialise path-link array. Each path is defined by a link utilisation array (one row in the path-link array). 1 indicates link corresponding to index is used, 0 indicates not used.

Parameters:

Name	Type	Description	Default
graph	Graph	NetworkX graph	required
k	int	Number of paths	required
disjoint	bool	Whether to use edge-disjoint paths. Defaults to False.	False
weight	str	Sort paths by edge attribute. Defaults to "weight".	'weight'
directed	bool	Whether graph is directed. Defaults to False.	False
modulations_array	Array	Array of maximum spectral efficiency for modulation format on path. Defaults to None.	None
rwa_lr	bool	Whether the environment is RWA with lightpath reuse (affects path ordering).	False
path_snr	bool	If GN model is used, include extra row of zeroes for unutilised paths	False

Returns:

Type	Description
Array	chex.Array: Path-link array (N(N-1)*k x E) where N is number of nodes, E is number of edges, k is number of shortest paths

Source code in xlron/environments/env_funcs.py

def init_path_link_array(
        graph: nx.Graph,
        k: int,
        disjoint: bool = False,
        weight: str = "weight",
        directed: bool = False,
        modulations_array: chex.Array = None,
        rwa_lr: bool = False,
        scale_factor: float = 1.0,
        path_snr: bool = False,
) -> chex.Array:
    """Initialise path-link array.
    Each path is defined by a link utilisation array (one row in the path-link array).
    1 indicates link corresponding to index is used, 0 indicates not used.

    Args:
        graph (nx.Graph): NetworkX graph
        k (int): Number of paths
        disjoint (bool, optional): Whether to use edge-disjoint paths. Defaults to False.
        weight (str, optional): Sort paths by edge attribute. Defaults to "weight".
        directed (bool, optional): Whether graph is directed. Defaults to False.
        modulations_array (chex.Array, optional): Array of maximum spectral efficiency for modulation format on path. Defaults to None.
        rwa_lr (bool, optional): Whether the environment is RWA with lightpath reuse (affects path ordering).
        path_snr (bool, optional): If GN model is used, include extra row of zeroes for unutilised paths
        to ensure correct SNR calculation for empty paths (path index -1).

    Returns:
        chex.Array: Path-link array (N(N-1)*k x E) where N is number of nodes, E is number of edges, k is number of shortest paths
    """
    def get_k_shortest_paths(g, source, target, k, weight):
        return list(
            islice(nx.shortest_simple_paths(g, source, target, weight=weight), k)
        )

    def get_k_disjoint_shortest_paths(g, source, target, k, weight):
        k_paths_disjoint_unsorted = list(nx.edge_disjoint_paths(g, source, target))
        k_paths_shortest = get_k_shortest_paths(g, source, target, k, weight=weight)

        # Keep disjoint paths and add unique shortest paths until k paths reached
        disjoint_ids = [tuple(path) for path in k_paths_disjoint_unsorted]
        k_paths = k_paths_disjoint_unsorted
        for path in k_paths_shortest:
            if tuple(path) not in disjoint_ids:
                k_paths.append(path)
        k_paths = k_paths[:k]
        return k_paths

    paths = []
    edges = sorted(graph.edges)

    # Get the k-shortest paths for each node pair
    k_path_collections = []
    get_paths = get_k_disjoint_shortest_paths if disjoint else get_k_shortest_paths
    for node_pair in combinations(graph.nodes, 2):

        k_paths = get_paths(graph, node_pair[0], node_pair[1], k, weight=weight)
        k_path_collections.append(k_paths)

    if directed:  # Get paths in reverse direction
        for node_pair in combinations(graph.nodes, 2):
            k_paths_rev = get_paths(graph, node_pair[1], node_pair[0], k, weight=weight)
            k_path_collections.append(k_paths_rev)

    # Sort the paths for each node pair
    max_missing_paths = 0
    for k_paths in k_path_collections:

        source, dest = k_paths[0][0], k_paths[0][-1]

        # Sort the paths by # of hops then by length, or just length
        path_lengths = [nx.path_weight(graph, path, weight='weight') for path in k_paths]
        path_num_links = [len(path) - 1 for path in k_paths]

        # Get maximum spectral efficiency for modulation format on path
        if modulations_array is not None and rwa_lr is not True:
            se_of_path = []
            modulations_array = modulations_array[::-1]
            for length in path_lengths:
                for modulation in modulations_array:
                    if length <= modulation[0]:
                        se_of_path.append(modulation[1])
                        break
            # Sorting by the num_links/se instead of just path length is observed to improve performance
            path_weighting = [num_links/se for se, num_links in zip(se_of_path, path_num_links)]
        elif rwa_lr:
            path_capacity = [float(calculate_path_capacity(path_length, scale_factor=scale_factor)) for path_length in path_lengths]
            path_weighting = [num_links/path_capacity for num_links, path_capacity in zip(path_num_links, path_capacity)]
        elif weight is None:
            path_weighting = path_num_links
        else:
            path_weighting = path_lengths

        # if less then k unique paths, add empty paths
        empty_path = [0] * len(graph.edges)
        num_missing_paths = k - len(k_paths)
        max_missing_paths = max(max_missing_paths, num_missing_paths)
        k_paths = k_paths + [empty_path] * num_missing_paths
        path_weighting = path_weighting + [1e6] * num_missing_paths
        path_lengths = path_lengths + [1e6] * num_missing_paths

        # Sort by number of links then by length (or just by length if weight is specified)
        unsorted_paths = zip(k_paths, path_weighting, path_lengths)
        k_paths_sorted = [(source, dest, weighting, path) for path, weighting, _ in sorted(unsorted_paths, key=lambda x: (x[1], 1/x[2]) if weight is None else x[2])]

        # Keep only first k paths
        k_paths_sorted = k_paths_sorted[:k]

        prev_link_usage = empty_path
        for k_path in k_paths_sorted:
            k_path = k_path[-1]
            link_usage = [0]*len(graph.edges)  # Initialise empty path
            if sum(k_path) == 0:
                link_usage = prev_link_usage
            else:
                for i in range(len(k_path)-1):
                    s, d = k_path[i], k_path[i + 1]
                    for edge_index, edge in enumerate(edges):
                        condition = (edge[0] == s and edge[1] == d) if directed else \
                            ((edge[0] == s and edge[1] == d) or (edge[0] == d and edge[1] == s))
                        if condition:
                            link_usage[edge_index] = 1
            path = link_usage
            prev_link_usage = link_usage
            paths.append(path)

    # If using GN model, add extra row of zeroes for empty paths for SNR calculation
    if path_snr:
        empty_path = [0] * len(graph.edges)
        paths.append(empty_path)

    return jnp.array(paths, dtype=SMALL_INT_DTYPE)

init_path_se_array(path_length_array, modulations_array)

Initialise array of maximum spectral efficiency for highest-order modulation format on path.

Parameters:

Name	Type	Description	Default
path_length_array	array	Array of path lengths	required
modulations_array	array	Array of maximum spectral efficiency for modulation format on path	required

Returns:

Type	Description
	jnp.array: Array of maximum spectral efficiency for on path

Source code in xlron/environments/env_funcs.py

def init_path_se_array(path_length_array, modulations_array):
    """Initialise array of maximum spectral efficiency for highest-order modulation format on path.

    Args:
        path_length_array (jnp.array): Array of path lengths
        modulations_array (jnp.array): Array of maximum spectral efficiency for modulation format on path

    Returns:
        jnp.array: Array of maximum spectral efficiency for on path
    """
    se_list = []
    # Flip the modulation array so that the shortest path length is first
    modulations_array = modulations_array[::-1]
    for length in path_length_array:
        for modulation in modulations_array:
            if length <= modulation[0]:
                se_list.append(modulation[1])
                break
    return jnp.array(se_list, dtype=SMALL_INT_DTYPE)

init_rsa_request_array()

Initialize request array

Source code in xlron/environments/env_funcs.py

def init_rsa_request_array():
    """Initialize request array"""
    return jnp.zeros(3, dtype=MED_INT_DTYPE)

init_traffic_matrix(key, params)

Initialize traffic matrix. Allows for random traffic matrix or uniform traffic matrix. Source-dest traffic requests are sampled probabilistically from the resulting traffic matrix.

Parameters:

Name	Type	Description	Default
key	PRNGKey	PRNG key	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
	jnp.array: Traffic matrix

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def init_traffic_matrix(key: chex.PRNGKey, params: EnvParams):
    """Initialize traffic matrix. Allows for random traffic matrix or uniform traffic matrix.
    Source-dest traffic requests are sampled probabilistically from the resulting traffic matrix.

    Args:
        key (chex.PRNGKey): PRNG key
        params (EnvParams): Environment parameters

    Returns:
        jnp.array: Traffic matrix
    """
    if params.random_traffic:
        traffic_matrix = jax.random.uniform(key, shape=(params.num_nodes, params.num_nodes), dtype=SMALL_FLOAT_DTYPE)
    else:
        traffic_matrix = jnp.ones((params.num_nodes, params.num_nodes), dtype=SMALL_FLOAT_DTYPE)
    diag_elements = jnp.diag_indices_from(traffic_matrix)
    # Set main diagonal to zero so no requests from node to itself
    traffic_matrix = traffic_matrix.at[diag_elements].set(0)
    traffic_matrix = normalise_traffic_matrix(traffic_matrix)
    return traffic_matrix

init_transceiver_amplifier_noise_arrays(link_resources, ref_lambda, slot_size, noise_data_filepath=None)

Initialise transceiver and amplifier noise arrays. Args: link_resources (int): Number of link resources ref_lambda (float): Reference wavelength slot_size (float): Slot size noise_data_filepath (str, optional): Path to CSV file containing modulation formats. Defaults to None. Returns: Tuple[chex.Array, chex.Array]: Transceiver noise array, Amplifier noise array

Source code in xlron/environments/env_funcs.py

def init_transceiver_amplifier_noise_arrays(
        link_resources: int,
        ref_lambda: float,
        slot_size: float,
        noise_data_filepath: str = None
) -> Tuple[chex.Array, chex.Array]:
    """Initialise transceiver and amplifier noise arrays.
    Args:
        link_resources (int): Number of link resources
        ref_lambda (float): Reference wavelength
        slot_size (float): Slot size
        noise_data_filepath (str, optional): Path to CSV file containing modulation formats. Defaults to None.
    Returns:
        Tuple[chex.Array, chex.Array]: Transceiver noise array, Amplifier noise array
    """
    f = pathlib.Path(noise_data_filepath) if noise_data_filepath else (
            pathlib.Path(__file__).parents[1].absolute() / "data" / "gn_model" / "transceiver_amplifier_data.csv")
    noise_data = np.genfromtxt(f, delimiter=',')
    # Drop empty first row (headers) and column (name)
    noise_data = noise_data[1:, 1:]
    # Columns are: wavelength_min_nm,wavelength_max_nm,frequency_min_ghz,frequency_max_ghz,NF_ASE_dB,SNR_TRX_dB
    frequency_min_ghz = noise_data[:, 2]
    frequency_max_ghz = noise_data[:, 3]
    amplifier_noise_db = noise_data[:, 4]  # NF_ASE_dB
    transceiver_snr_db = noise_data[:, 5]  # SNR_TRX_dB

    # Define slot centres in GHz relative to central wavelength
    slot_centres = (jnp.arange(link_resources) - (link_resources - 1) / 2) * slot_size

    # Transform relative slot centres to absolute frequencies in GHz
    ref_frequency_ghz = c / ref_lambda / 1e9
    slot_frequencies_ghz = ref_frequency_ghz + slot_centres

    # Initialize output arrays
    transceiver_snr_array = jnp.zeros(link_resources)
    amplifier_noise_figure_array = jnp.zeros(link_resources)

    # For each slot, find which band it belongs to
    for i, freq in enumerate(slot_frequencies_ghz):
        # Find the band this frequency falls into
        for j in range(len(frequency_min_ghz)):
            if frequency_min_ghz[j] <= freq <= frequency_max_ghz[j]:
                transceiver_snr_array = transceiver_snr_array.at[i].set(transceiver_snr_db[j])
                amplifier_noise_figure_array = amplifier_noise_figure_array.at[i].set(amplifier_noise_db[j])
                break
        else:
            # If frequency is outside all bands, could raise error or use default
            raise ValueError(f"Frequency {freq} GHz is outside the defined bands")

    return transceiver_snr_array, amplifier_noise_figure_array

init_virtual_topology_patterns(pattern_names)

Initialise virtual topology patterns. First 3 digits comprise the "action counter": first index is num unique nodes, second index is total steps, final is remaining steps until completion of request. Remaining digits define the topology pattern, with 1 to indicate links and other positive integers are node indices.

Parameters:

Name	Type	Description	Default
pattern_names	list	List of virtual topology pattern names	required

Returns:

Type	Description
Array	chex.Array: Array of virtual topology patterns

Source code in xlron/environments/env_funcs.py

def init_virtual_topology_patterns(pattern_names: str) -> chex.Array:
    """Initialise virtual topology patterns.
    First 3 digits comprise the "action counter": first index is num unique nodes, second index is total steps,
    final is remaining steps until completion of request.
    Remaining digits define the topology pattern, with 1 to indicate links and other positive integers are node indices.

    Args:
        pattern_names (list): List of virtual topology pattern names

    Returns:
        chex.Array: Array of virtual topology patterns
    """
    patterns = []
    # TODO - Allow 2 node requests in VONE (check if any modifications necessary other than below)
    #if "2_bus" in pattern_names:
    #    patterns.append([2, 1, 1, 2, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0])
    if "3_bus" in pattern_names:
        patterns.append([3, 2, 2, 2, 1, 3, 1, 4])
    if "3_ring" in pattern_names:
        patterns.append([3, 3, 3, 2, 1, 3, 1, 4, 1, 2])
    if "4_bus" in pattern_names:
        patterns.append([4, 3, 3, 2, 1, 3, 1, 4, 1, 5])
    if "4_ring" in pattern_names:
        patterns.append([4, 4, 4, 2, 1, 3, 1, 4, 1, 5, 1, 2])
    if "5_bus" in pattern_names:
        patterns.append([5, 4, 4, 2, 1, 3, 1, 4, 1, 5, 1, 6])
    if "5_ring" in pattern_names:
        patterns.append([5, 5, 5, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 2])
    if "6_bus" in pattern_names:
        patterns.append([6, 5, 5, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 7])
    max_length = max([len(pattern) for pattern in patterns])
    # Pad patterns with zeroes to match longest
    for pattern in patterns:
        pattern.extend([0]*(max_length-len(pattern)))
    return jnp.array(patterns, dtype=SMALL_INT_DTYPE)

init_vone_request_array(params)

Initialize request array either with uniform resources

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(0,))
def init_vone_request_array(params: EnvParams):
    """Initialize request array either with uniform resources"""
    return jnp.zeros((2, params.max_edges*2+1,), dtype=MED_INT_DTYPE)

kca_ff(state, params)

Congestion-aware First Fit. Only suitable for RSA/RMSA. Method:

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def kca_ff(state: EnvState, params: EnvParams) -> chex.Array:
    """Congestion-aware First Fit. Only suitable for RSA/RMSA.
    Method:

    """
    mask = get_action_mask(state, params)
    # Get index of first available slots for each path
    first_slots = first_fit(state, params)
    # Get nodes
    nodes_sd, _ = read_rsa_request(state.request_array)
    # Initialise array to hold congestion on each path
    path_congestion_array = jnp.full((mask.shape[0],), 0.)
    link_weights = get_link_weights(state, params)

    def get_path_congestion(i, val):
        # Get links on path
        path = get_paths(params, nodes_sd)[i]
        # Get congestion
        path_link_congestion = jnp.multiply(link_weights, path)
        path_congestion = jnp.sum(path_link_congestion).reshape((1,))
        return jax.lax.dynamic_update_slice(val, path_congestion, (i,))

    path_congestion_array = jax.lax.fori_loop(0, mask.shape[0], get_path_congestion, path_congestion_array)
    path_index = jnp.argmin(path_congestion_array)
    slot_index = first_slots[path_index] % params.link_resources
    action = path_index * params.link_resources + slot_index
    return action

kmc_ff(state, params)

K-Minimum Cut. Only suitable for RSA/RMSA. Method: 1. Go through action mask and find the first available slot on all paths. 2. For each path, allocate the first available slot. 3. Sum number of new consecutive zero regions (cuts) created by assignment (on each link) 4. Choose path that creates the fewest cuts.

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def kmc_ff(state: EnvState, params: EnvParams) -> chex.Array:
    """K-Minimum Cut. Only suitable for RSA/RMSA.
    Method:
    1. Go through action mask and find the first available slot on all paths.
    2. For each path, allocate the first available slot.
    3. Sum number of new consecutive zero regions (cuts) created by assignment (on each link)
    4. Choose path that creates the fewest cuts.
    """
    mask = get_action_mask(state, params)
    first_slots = first_fit(state, params)
    link_slot_array = jnp.where(state.link_slot_array < 0, 1., state.link_slot_array)
    nodes_sd, requested_bw = read_rsa_request(state.request_array)
    block_sizes = jax.vmap(find_block_sizes, in_axes=(0,))(link_slot_array)
    block_sizes_mask = jnp.where(block_sizes > 0, 1, 0.)  # Binary array showing initial block starts
    block_count = jnp.sum(block_sizes_mask, axis=1)

    def get_cuts_on_path(i, result):
        initial_slot_index = first_slots[i] % params.link_resources
        path = get_paths(params, nodes_sd)[i]
        se = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else 1
        num_slots = required_slots(requested_bw, se, params.slot_size, guardband=params.guardband)
        # Make link-slot_array positive
        updated_slots = vmap_set_path_links(link_slot_array, path, initial_slot_index, num_slots, 1.)
        updated_block_sizes = jax.vmap(find_block_sizes, in_axes=(0,))(updated_slots)
        updated_block_sizes_mask = jnp.where(updated_block_sizes > 0, 1, 0)  # Binary array showing updated block starts
        updated_block_count = jnp.sum(updated_block_sizes_mask, axis=1)
        num_cuts = jax.lax.cond(
            mask[i][initial_slot_index] == 0.,  # If true, no valid action for path
            lambda x: jnp.full((1,), params.link_resources*params.num_links).astype(jnp.float32),  # Return max no. of cuts
            lambda x: jnp.sum(jnp.maximum(updated_block_count - block_count, 0.)).reshape((1,)),  # Else, return number of cuts
            1.
        )
        result = jax.lax.dynamic_update_slice(result, num_cuts, (i,))
        return result

    # Initialise array to hold number of cuts on each path
    path_cuts_array = jnp.full((mask.shape[0],), 0.)
    path_cuts_array = jax.lax.fori_loop(0, mask.shape[0], get_cuts_on_path, path_cuts_array)
    path_index = jnp.argmin(path_cuts_array)
    slot_index = first_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

kme_ff(state, params)

K-Minimum Entropy. Only suitable for RSA/RMSA. Method: 1. Go through action mask and find the first available slot on all paths. 2. For each path, allocate the first available slot. 3. Sum number of new consecutive zero regions (cuts) created by assignment (on each link) 4. Choose path that creates the fewest cuts.

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def kme_ff(state: EnvState, params: EnvParams) -> chex.Array:
    """K-Minimum Entropy. Only suitable for RSA/RMSA.
    Method:
    1. Go through action mask and find the first available slot on all paths.
    2. For each path, allocate the first available slot.
    3. Sum number of new consecutive zero regions (cuts) created by assignment (on each link)
    4. Choose path that creates the fewest cuts.
    """
    mask = get_action_mask(state, params)
    first_slots = first_fit(state, params)
    link_slot_array = jnp.where(state.link_slot_array < 0, 1., state.link_slot_array)
    nodes_sd, requested_bw = read_rsa_request(state.request_array)
    max_entropy = jnp.sum(jnp.log(params.link_resources)) * params.num_links

    def get_link_entropy(blocks):
        ent = jax.vmap(lambda x: jnp.sum(x/params.link_resources * jnp.log(params.link_resources/x)), in_axes=0)(blocks)
        return jnp.sum(jnp.where(blocks > 0, ent, 0))

    def get_entropy_on_path(i, result):
        initial_slot_index = first_slots[i] % params.link_resources
        path = get_paths(params, nodes_sd)[i]
        se = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else 1
        num_slots = required_slots(requested_bw, se, params.slot_size, guardband=params.guardband)
        # Make link-slot_array positive
        updated_slots = vmap_set_path_links(link_slot_array, path, initial_slot_index, num_slots, 1.)
        updated_block_sizes = jax.vmap(find_block_sizes, in_axes=(0,))(updated_slots)
        updated_entropy = jax.vmap(get_link_entropy, in_axes=(0,))(updated_block_sizes)
        new_path_entropy = jnp.sum(jnp.dot(path, updated_entropy)).reshape((1,))
        new_path_entropy = jax.lax.cond(
            mask[i][initial_slot_index] == 0.,  # If true, no valid action for path
            lambda x: max_entropy.astype(jnp.float32).reshape((1,)),  # Return maximum entropy
            lambda x: new_path_entropy,  # Else, return number of cuts
            1.
        )
        result = jax.lax.dynamic_update_slice(result, new_path_entropy, (i,))
        return result

    path_entropy_array = jnp.full((mask.shape[0],), 0.)
    path_entropy_array = jax.lax.fori_loop(0, mask.shape[0], get_entropy_on_path, path_entropy_array)
    path_index = jnp.argmin(path_entropy_array)
    slot_index = first_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

kmf_ff(state, params)

K-Minimum Frag-size. Only suitable for RSA/RMSA. Method: 1. Go through action mask and find the first available slot on all paths. 2. For each path, allocate the first available slot. 3. Sum number of new consecutive zero regions (cuts) created by assignment (on each link) 4. Choose path that creates the fewest cuts.

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def kmf_ff(state: RSAEnvState, params: RSAEnvParams) -> chex.Array:
    """K-Minimum Frag-size. Only suitable for RSA/RMSA.
    Method:
    1. Go through action mask and find the first available slot on all paths.
    2. For each path, allocate the first available slot.
    3. Sum number of new consecutive zero regions (cuts) created by assignment (on each link)
    4. Choose path that creates the fewest cuts.
    """
    mask = get_action_mask(state, params)
    first_slots = first_fit(state, params)
    link_slot_array = jnp.where(state.link_slot_array < 0, 1., state.link_slot_array)
    nodes_sd, requested_bw = read_rsa_request(state.request_array)
    blocks = jax.vmap(find_block_sizes, in_axes=(0,))(link_slot_array)

    def get_frags_on_path(i, result):
        initial_slot_index = first_slots[i] % params.link_resources
        path = get_paths(params, nodes_sd)[i]
        se = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else 1
        num_slots = required_slots(requested_bw, se, params.slot_size, guardband=params.guardband)
        # Mask on path links
        block_sizes = jax.vmap(lambda x, y: jnp.where(x > 0, y, 0.), in_axes=(0, 0))(path, blocks)
        updated_slots = vmap_set_path_links(state.link_slot_array, path, initial_slot_index, num_slots, -1)
        updated_block_sizes = jax.vmap(find_block_sizes, in_axes=(0,))(updated_slots)
        # Mask on path links
        updated_block_sizes = jax.vmap(lambda x, y: jnp.where(x > 0, y, 0.), in_axes=(0, 0))(path, updated_block_sizes)
        difference = updated_block_sizes - block_sizes
        new_frags = jnp.where(difference != 0, block_sizes + difference, 0.)
        # Slice new frags up to initial slot index (so as to only consider frags to the left)
        new_frags = jnp.where(jnp.arange(params.link_resources) < initial_slot_index, new_frags, 0.)
        new_frag_size = jnp.sum(new_frags)
        num_frags = jax.lax.cond(
            mask[i][initial_slot_index] == 0.,  # If true, no valid action for path
            lambda x: jnp.full((1,), float(params.link_resources * params.num_links)),  # Return max frag size
            lambda x: new_frag_size.reshape((1,)),
            # Else, return number of cuts
            1.
        )
        result = jax.lax.dynamic_update_slice(result, num_frags, (i,))
        return result

    # Initialise array to hold number of cuts on each path
    path_frags_array = jnp.full((mask.shape[0],), 0.)
    path_frags_array = jax.lax.fori_loop(0, mask.shape[0], get_frags_on_path, path_frags_array)
    path_index = jnp.argmin(path_frags_array)
    slot_index = first_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

ksp_bf(state, params)

Get the first available slot from all k-shortest paths Method: Go through action mask and find the first available slot, starting from shortest path

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def ksp_bf(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the first available slot from all k-shortest paths
    Method: Go through action mask and find the first available slot, starting from shortest path

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Action
    """
    best_slots, fitness = best_fit(state, params)
    # Chosen path is the first one with an available slot
    path_index = jnp.argmin(jnp.where(fitness < jnp.inf, 0, 1))
    slot_index = best_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

ksp_ff(state, params)

Get the first available slot from the shortest available path Method: Go through action mask and find the first available slot, starting from shortest path

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def ksp_ff(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the first available slot from the shortest available path
    Method: Go through action mask and find the first available slot, starting from shortest path

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Action
    """
    first_slots = first_fit(state, params)
    # Chosen path is the first one with an available slot
    path_index = jnp.argmax(first_slots < params.link_resources)
    slot_index = first_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

ksp_ff_multiband(state, params)

Get the first available slot from all k-shortest paths in multiband scenario Method: Go through action mask and find the first available slot, starting from shortest path

Parameters:

Name	Type	Description	Default
state	MultiBandRSAEnvState	Environment state specific to multiband operations	required
params	MultiBandRSAEnvParams	Environment parameters including multiband details	required

Returns: chex.Array: Action

Source code in xlron/heuristics/heuristics.py

def ksp_ff_multiband(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the first available slot from all k-shortest paths in multiband scenario
    Method: Go through action mask and find the first available slot, starting from shortest path

    Args:
        state (MultiBandRSAEnvState): Environment state specific to multiband operations
        params (MultiBandRSAEnvParams): Environment parameters including multiband details
    Returns:
        chex.Array: Action
    """
    pass

ksp_lf(state, params)

Get the last available slot on the shortest available path Method: Go through action mask and find the last available slot, starting from shortest path

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def ksp_lf(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the last available slot on the shortest available path
    Method: Go through action mask and find the last available slot, starting from shortest path

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Action
    """
    last_slots = last_fit(state, params)
    # Chosen path is the first one with an available slot
    path_index = jnp.argmax(last_slots < params.link_resources)
    slot_index = last_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

ksp_mu(state, params, unique_lightpaths, relative)

Get the most-used slot on the shortest available path. Method: Go through action mask and find the utilisation of available slots on each path. Find the shortest available path and choose the most utilised slot on that path.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
unique_lightpaths	bool	Whether to consider unique lightpaths	required
relative	bool	Whether to return relative utilisation	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1, 2, 3))
def ksp_mu(state: EnvState, params: EnvParams, unique_lightpaths: bool, relative: bool) -> chex.Array:
    """Get the most-used slot on the shortest available path.
    Method: Go through action mask and find the utilisation of available slots on each path.
    Find the shortest available path and choose the most utilised slot on that path.

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        unique_lightpaths (bool): Whether to consider unique lightpaths
        relative (bool): Whether to return relative utilisation

    Returns:
        chex.Array: Action
    """
    mask = get_action_mask(state, params)
    most_used_slots = most_used(state, params, unique_lightpaths, relative)
    # Get usage of available slots
    most_used_mask = most_used_slots * mask
    # Get index of most-used available slot for each path
    most_used_slots = jnp.argmax(most_used_mask, axis=1).astype(jnp.int32)
    # Chosen path is the first one with an available slot
    available_paths = jnp.max(mask, axis=1)
    path_index = jnp.argmax(available_paths)
    slot_index = most_used_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

last_fit(state, params)

Last-Fit Spectrum Allocation. Returns the last fit slot for each path.

Source code in xlron/heuristics/heuristics.py

def last_fit(state: EnvState, params: EnvParams) -> chex.Array:
    """Last-Fit Spectrum Allocation. Returns the last fit slot for each path."""
    mask = get_action_mask(state, params)
    # Add a column of ones to the mask to make sure that occupied paths have non-zero index in "last_slots"
    mask = jnp.concatenate((jnp.full((mask.shape[0], 1), 1), mask), axis=1)
    # Get index of last available slots for each path
    last_slots = jnp.argmax(mask[:, ::-1], axis=1)
    # Convert to index from the left
    last_slots = params.link_resources - last_slots - 1
    return last_slots

lf_ksp(state, params)

Get the last available slot from all paths Method: Go through action mask and find the last available slot on all paths

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1,))
def lf_ksp(state: EnvState, params: EnvParams) -> chex.Array:
    """Get the last available slot from all paths
    Method: Go through action mask and find the last available slot on all paths

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters

    Returns:
        chex.Array: Action
    """
    last_slots = last_fit(state, params)
    # Chosen path is the one with the highest index of last available slot
    path_index = jnp.argmax(last_slots)
    slot_index = last_slots[path_index] % params.link_resources
    # Convert indices to action
    action = path_index * params.link_resources + slot_index
    return action

make_graph(topology_name='conus', topology_directory=None)

Create graph from topology definition. Topologies must be defined in JSON format in the topologies directory and named as the topology name with .json extension.

Parameters:

Name	Type	Description	Default
topology_name	str	topology name	'conus'
topology_directory	str	topology directory	None

Returns:

Name	Type	Description
graph		graph

Source code in xlron/environments/env_funcs.py

def make_graph(topology_name: str = "conus", topology_directory: str = None):
    """Create graph from topology definition.
    Topologies must be defined in JSON format in the topologies directory and
    named as the topology name with .json extension.

    Args:
        topology_name: topology name
        topology_directory: topology directory

    Returns:
        graph: graph
    """
    topology_path = pathlib.Path(topology_directory) if topology_directory else (
            pathlib.Path(__file__).parents[1].absolute() / "data" / "topologies")
    # Create topology
    if topology_name == "4node":
        # 4 node ring
        graph = nx.from_numpy_array(np.array([[0, 1, 0, 1],
                                            [1, 0, 1, 0],
                                               [0, 1, 0, 1],
                                               [1, 0, 1, 0]]))
        # Add edge weights to graph
        nx.set_edge_attributes(graph, {(0, 1): 4, (1, 2): 3, (2, 3): 2, (3, 0): 1}, "weight")
    elif topology_name == "7node":
        # 7 node ring
        graph = nx.from_numpy_array(jnp.array([[0, 1, 0, 0, 0, 0, 1],
                                               [1, 0, 1, 0, 0, 0, 0],
                                               [0, 1, 0, 1, 0, 0, 0],
                                               [0, 0, 1, 0, 1, 0, 0],
                                               [0, 0, 0, 1, 0, 1, 0],
                                               [0, 0, 0, 0, 1, 0, 1],
                                               [1, 0, 0, 0, 0, 1, 0]]))
        # Add edge weights to graph
        nx.set_edge_attributes(graph, {(0, 1): 4, (1, 2): 3, (2, 3): 2, (3, 4): 1, (4, 5): 2, (5, 6): 3, (6, 0): 4}, "weight")
    else:
        with open(topology_path / f"{topology_name}.json") as f:
            graph = nx.node_link_graph(json.load(f))
    return graph

mask_nodes(state, num_nodes)

Returns mask of valid actions for node selection. 1 for valid action, 0 for invalid action.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
num_nodes	Scalar	Number of nodes	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_nodes(state: EnvState, num_nodes: chex.Scalar) -> EnvState:
    """Returns mask of valid actions for node selection. 1 for valid action, 0 for invalid action.

    Args:
        state: Environment state
        num_nodes: Number of nodes

    Returns:
        state: Updated environment state
    """
    total_actions = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.action_counter, 1, 1))
    remaining_actions = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.action_counter, 2, 1))
    full_request = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.request_array, 0, 1))
    virtual_topology = jnp.squeeze(jax.lax.dynamic_slice_in_dim(state.request_array, 1, 1))
    request = jax.lax.dynamic_slice_in_dim(full_request, (remaining_actions - 1) * 2, 3)
    node_request_s = jax.lax.dynamic_slice_in_dim(request, 2, 1)
    node_request_d = jax.lax.dynamic_slice_in_dim(request, 0, 1)
    prev_action = jax.lax.dynamic_slice_in_dim(state.action_history, (remaining_actions) * 2, 3)
    prev_dest = jax.lax.dynamic_slice_in_dim(prev_action, 0, 1)
    node_indices = jnp.arange(0, num_nodes)
    # Get requested indices from request array virtual topology
    requested_indices = jax.lax.dynamic_slice_in_dim(virtual_topology, (remaining_actions-1)*2, 3)
    requested_index_d = jax.lax.dynamic_slice_in_dim(requested_indices, 0, 1)
    # Get index of previous selected node
    prev_selected_node = jnp.where(virtual_topology == requested_index_d, state.action_history, jnp.full(virtual_topology.shape, -1))
    # will be current index if node only occurs once in virtual topology or will be different index if occurs more than once
    prev_selected_index = jnp.argmax(prev_selected_node).astype(MED_INT_DTYPE)
    prev_selected_node_d = jax.lax.dynamic_slice_in_dim(state.action_history, prev_selected_index, 1)

    # If first action, source and dest both to be assigned -> just mask all nodes based on resources
    # Thereafter, source must be previous dest. Dest can be any node (except previous allocations).
    state = state.replace(
        node_mask_s=jax.lax.cond(
            jnp.equal(remaining_actions, total_actions),
            lambda x: jnp.where(
                state.node_capacity_array >= node_request_s,
                x,
                jnp.zeros(num_nodes)
            ),
            lambda x: jnp.where(
                node_indices == prev_dest,
                x,
                jnp.zeros(num_nodes)
            ),
            jnp.ones(num_nodes),
        )
    )
    state = state.replace(
        node_mask_d=jnp.where(
            state.node_capacity_array >= node_request_d,
            jnp.ones(num_nodes),
            jnp.zeros(num_nodes)
        )
    )
    # If not first move, set node_mask_d to zero wherever node_mask_s is 1
    # to avoid same node selection for s and d
    state = state.replace(
        node_mask_d=jax.lax.cond(
            jnp.equal(remaining_actions, total_actions),
            lambda x: x,
            lambda x: jnp.where(
                state.node_mask_s == 1,
                jnp.zeros(num_nodes),
                x
            ),
            state.node_mask_d,
        )
    )

    def mask_previous_selections(i, val):
        # Disallow previously allocated nodes
        update_slice = lambda j, x: jax.lax.dynamic_update_slice_in_dim(x, jnp.array([0.]), j, axis=0)
        val = jax.lax.cond(
            i % 2 == 0,
            lambda x: update_slice(x[0][i], x[1]),  # i is node request index
            lambda x: update_slice(x[0][i+1], x[1]),  # i is slot request index (so add 1 to get next node)
            (state.action_history, val),
        )
        return val

    state = state.replace(
        node_mask_d=jax.lax.fori_loop(
            remaining_actions*2,
            state.action_history.shape[0]-1,
            mask_previous_selections,
            state.node_mask_d
        )
    )
    # If requested node index is new then disallow previously allocated nodes
    # If not new, then must match previously allocated node for that index
    state = state.replace(
        node_mask_d=jax.lax.cond(
            jnp.squeeze(prev_selected_node_d) >= 0,
            lambda x: jnp.where(
                node_indices == prev_selected_node_d,
                x[1],
                x[0],
            ),
            lambda x: x[2],
            (jnp.zeros(num_nodes), jnp.ones(num_nodes), state.node_mask_d),
        )
    )
    return state

mask_slots(state, params, request)

Returns binary mask of valid actions. 1 for valid action, 0 for invalid action.

1. Check request for source and destination nodes
2. For each path:
   • Get current slots on path (with padding on end to avoid out of bounds)
   • Get mask for required slots on path
   • Multiply through current slots with required slots mask to check if slots available on path
   • Remove padding from mask
   • Return path mask
3. Update total mask with path mask
4. If aggregate_slots > 1, aggregate slot mask to reduce action space

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_slots(state: EnvState, params: EnvParams, request: chex.Array) -> EnvState:
    """Returns binary mask of valid actions. 1 for valid action, 0 for invalid action.

    1. Check request for source and destination nodes
    2. For each path:
        - Get current slots on path (with padding on end to avoid out of bounds)
        - Get mask for required slots on path
        - Multiply through current slots with required slots mask to check if slots available on path
        - Remove padding from mask
        - Return path mask
    3. Update total mask with path mask
    4. If aggregate_slots > 1, aggregate slot mask to reduce action space

    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_mask = jnp.zeros((params.link_resources * params.k_paths), dtype=LARGE_FLOAT_DTYPE)

    def mask_path(i, mask):
        # Get slots for path
        slots = get_path_slots(state.link_slot_array, params, nodes_sd, i)
        # Add padding to slots at end
        slots = jnp.concatenate((slots, jnp.ones(params.max_slots, dtype=LARGE_FLOAT_DTYPE)))
        # Convert bandwidth to slots for each path
        spectral_efficiency = get_paths_se(params, nodes_sd)[i] if params.consider_modulation_format else one
        requested_slots = required_slots(requested_datarate, spectral_efficiency, params.slot_size, guardband=params.guardband)
        # Get mask used to check if request will fit slots
        request_mask = get_request_mask(requested_slots[0], params)

        def check_slots_available(j, val):
            # Multiply through by request mask to check if slots available
            request_slice = jax.lax.dynamic_slice(val, (j,), (params.max_slots,))
            slot_sum = jnp.sum(request_mask * request_slice, promote_integers=False) <= zero
            slot_sum = slot_sum.reshape((1,)).astype(LARGE_FLOAT_DTYPE)
            return jax.lax.dynamic_update_slice(val, slot_sum, (j,))

        # Mask out slots that are not valid
        path_mask = jax.lax.fori_loop(
            0,
            int(params.link_resources+1),  # No need to check last requested_slots-1 slots
            check_slots_available,
            slots,
        )
        # Cut off padding
        path_mask = jax.lax.dynamic_slice(path_mask, (0,), (params.link_resources,))
        # Update total mask with path mask
        mask = jax.lax.dynamic_update_slice(mask, path_mask, (i * params.link_resources,))
        return mask

    # Loop over each path
    link_slot_mask = jax.lax.fori_loop(0, params.k_paths, mask_path, init_mask)
    if params.aggregate_slots > 1:
        # Full link slot mask is used in process_path_action to get the correct slot from the aggregated slot action
        state = state.replace(full_link_slot_mask=link_slot_mask)
        link_slot_mask, _ = aggregate_slots(link_slot_mask.reshape(params.k_paths, -1), params)
        link_slot_mask = link_slot_mask.reshape(-1)
    state = state.replace(link_slot_mask=link_slot_mask)
    return state

mask_slots_rmsa_gn_model(state, params, request)

For use in RSAGNModelEnv. 1. For each path: 1.1 Get path slots 1.2 Get launch power

Parameters:

Name	Type	Description	Default
state	RSAGNModelEnvState	Environment state	required
params	RSAGNModelEnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_slots_rmsa_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams, request: chex.Array) -> EnvState:
    """For use in RSAGNModelEnv.
    1. For each path:
        1.1 Get path slots
        1.2 Get launch power


    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_mask = jnp.zeros((params.link_resources * params.k_paths))

    def mask_path(i, mask):
        path = get_paths(params, nodes_sd)[i]
        # Get slots for path
        slots = get_path_slots(state.link_slot_array, params, nodes_sd, i)
        # Add padding to slots at end
        # 0 means slot is free, 1 is occupied
        slots = jnp.concatenate((slots, jnp.ones(params.max_slots)))
        launch_power = get_launch_power(state, i, state.launch_power_array[i], params)
        lightpath_index = get_lightpath_index(params, nodes_sd, i)

        # This function checks through each available modulation format, checks the first and last available slots,
        # calculates the SNR, checks it meets the requirements, and returns the resulting mask
        def check_modulation_format(mod_format_index, init_path_mask):
            se = params.modulations_array.val[mod_format_index][1]
            req_slots = required_slots(requested_datarate, se, params.slot_size, guardband=params.guardband)[0]
            bandwidth_per_subchannel = params.slot_size
            req_snr = params.modulations_array.val[mod_format_index][2] + params.snr_margin
            # Get mask used to check if request will fit slots
            request_mask = get_request_mask(req_slots, params)

            def check_slots_available(j, val):
                # Multiply through by request mask to check if slots available
                slot_sum = jnp.sum(request_mask * jax.lax.dynamic_slice(val, (j,), (params.max_slots,)), promote_integers=False) <= 0
                slot_sum = slot_sum.reshape((1,)).astype(LARGE_FLOAT_DTYPE)
                return jax.lax.dynamic_update_slice(val, slot_sum, (j,))

            # Mask out slots that are not valid
            slot_mask = jax.lax.fori_loop(
                0,
                int(params.link_resources + 1),  # No need to check last requested_slots-1 slots
                check_slots_available,
                slots,
            )
            # Cut off padding
            slot_mask = jax.lax.dynamic_slice(slot_mask, (0,), (params.link_resources,))
            # Check first and last available slots for suitability
            ff_path_mask = jnp.concatenate((slot_mask, jnp.ones((1,))), axis=0)
            lf_path_mask = jnp.concatenate((jnp.ones((1,)), slot_mask), axis=0)
            first_available_slot_index = jnp.argmax(ff_path_mask)
            last_available_slot_index = params.link_resources - jnp.argmax(jnp.flip(lf_path_mask)) - 1
            # Assign "req_slots" subchannels (each with bandwidth = slot width) for the first and last possible slots
            ff_temp_state = state.replace(
                channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, first_available_slot_index, req_slots, bandwidth_per_subchannel),
                channel_power_array=vmap_set_path_links(state.channel_power_array, path, first_available_slot_index, req_slots, launch_power),
                path_index_array=vmap_set_path_links(state.path_index_array, path, first_available_slot_index, req_slots, lightpath_index),
                modulation_format_index_array=vmap_set_path_links(state.modulation_format_index_array, path, first_available_slot_index, req_slots, mod_format_index),
            )
            lf_temp_state = state.replace(
                channel_centre_bw_array=vmap_set_path_links(state.channel_centre_bw_array, path, last_available_slot_index, req_slots, bandwidth_per_subchannel),
                channel_power_array=vmap_set_path_links(state.channel_power_array, path, last_available_slot_index, req_slots, launch_power),
                path_index_array=vmap_set_path_links(state.path_index_array, path, last_available_slot_index, req_slots, lightpath_index),
                modulation_format_index_array=vmap_set_path_links(state.modulation_format_index_array, path, last_available_slot_index, req_slots, mod_format_index),
            )
            ff_temp_state = ff_temp_state.replace(link_snr_array=get_snr_link_array(ff_temp_state, params))
            lf_temp_state = lf_temp_state.replace(link_snr_array=get_snr_link_array(lf_temp_state, params))
            # Take the minimum value of SNR from all the subchannels
            ff_snr_value = get_minimum_snr_of_channels_on_path(
                ff_temp_state, path, first_available_slot_index, req_slots, params
            )
            lf_snr_value = get_minimum_snr_of_channels_on_path(
                lf_temp_state, path, last_available_slot_index, req_slots, params
            )
            # Check that other paths SNR is still sufficient (True if failure)
            ff_snr_check = 1 - check_action_rmsa_gn_model(ff_temp_state, None, params)
            lf_snr_check = 1 - check_action_rmsa_gn_model(lf_temp_state, None, params)
            ff_check = (ff_snr_value >= req_snr) * ff_snr_check
            lf_check = (lf_snr_value >= req_snr) * lf_snr_check

            slot_indices = jnp.arange(params.link_resources, dtype=MED_INT_DTYPE)
            mod_format_mask = jnp.where(slot_indices == first_available_slot_index, ff_check, False)
            mod_format_mask = jnp.where(slot_indices == last_available_slot_index, lf_check, mod_format_mask)
            path_mask = jnp.where(mod_format_mask, mod_format_index, init_path_mask)
            # jax.debug.print("ff_snr_check {}", ff_snr_check, ordered=True)
            # jax.debug.print("lf_snr_check {}", lf_snr_check, ordered=True)
            # jax.debug.print("ff_snr_value {}", ff_snr_value, ordered=True)
            # jax.debug.print("lf_snr_value {}", lf_snr_value, ordered=True)
            # jax.debug.print("first_available_slot_index {}", first_available_slot_index, ordered=True)
            # jax.debug.print("last_available_slot_index {}", last_available_slot_index, ordered=True)
            # jax.debug.print("req_snr {}", req_snr, ordered=True)
            # jax.debug.print("mod_format_mask {}", mod_format_mask, ordered=True)
            # jax.debug.print("path_mask {}", path_mask, ordered=True)
            return path_mask

        path_mask = jax.lax.fori_loop(0, params.modulations_array.val.shape[0], check_modulation_format, jnp.full((params.link_resources,), -1., dtype=LARGE_FLOAT_DTYPE))

        # Update total mask with path mask
        mask = jax.lax.dynamic_update_slice(mask, path_mask, (i * params.link_resources,))
        return mask

    # Loop over each path
    mod_format_mask = jax.lax.fori_loop(0, params.k_paths, mask_path, init_mask)
    link_slot_mask = jnp.where(mod_format_mask >= 0, 1.0, 0.0)
    if params.aggregate_slots > 1:
        # Full link slot mask is used in process_path_action to get the correct slot from the aggregated slot action
        state = state.replace(full_link_slot_mask=link_slot_mask)
        link_slot_mask, _ = aggregate_slots(link_slot_mask.reshape(params.k_paths, -1), params)
        link_slot_mask = link_slot_mask.reshape(-1)
    state = state.replace(
        link_slot_mask=link_slot_mask,
        mod_format_mask=mod_format_mask,
    )
    return state

mask_slots_rwalr(state, params, request)

For use in RWALightpathReuseEnv. Each lightpath has a maximum capacity defined in path_capacity_array. This is updated when a lightpath is assigned. If remaining path capacity is less than current request, corresponding link-slots are masked out. If link-slot is in use by another lightpath for a different source and destination node (even if not full) it is masked out. Step 1: - Mask out slots that are not valid based on path capacity (check link_capacity_array) Step 2: - Mask out slots that are not valid based on lightpath reuse (check path_index_array)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
request	Array	Request array in format [source_node, data-rate, destination_node]	required

Returns:

Name	Type	Description
state	EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def mask_slots_rwalr(state: EnvState, params: EnvParams, request: chex.Array) -> EnvState:
    """For use in RWALightpathReuseEnv.
    Each lightpath has a maximum capacity defined in path_capacity_array. This is updated when a lightpath is assigned.
    If remaining path capacity is less than current request, corresponding link-slots are masked out.
    If link-slot is in use by another lightpath for a different source and destination node (even if not full) it is masked out.
    Step 1:
    - Mask out slots that are not valid based on path capacity (check link_capacity_array)
    Step 2:
    - Mask out slots that are not valid based on lightpath reuse (check path_index_array)

    Args:
        state: Environment state
        params: Environment parameters
        request: Request array in format [source_node, data-rate, destination_node]

    Returns:
        state: Updated environment state
    """
    nodes_sd, requested_datarate = read_rsa_request(request)
    init_mask = jnp.zeros((params.link_resources * params.k_paths))
    source, dest = nodes_sd
    path_start_index = get_path_indices(source, dest, params.k_paths, params.num_nodes, directed=params.directed_graph)
    #jax.debug.print("path_start_index {}", path_start_index, ordered=True)
    #jax.debug.print("link_capacity_array {}", state.link_capacity_array, ordered=True)

    def mask_path(i, mask):
        # Step 1 - mask capacity
        capacity_mask = jnp.where(state.link_capacity_array < requested_datarate, 1., 0.)
        #jax.debug.print("capacity_mask {}", capacity_mask, ordered=True)
        capacity_slots = get_path_slots(capacity_mask, params, nodes_sd, i)
        #jax.debug.print("capacity_slots {}", capacity_slots, ordered=True)
        # Step 2 - mask lightpath reuse
        lightpath_index = path_start_index + i
        #jax.debug.print("lightpath_index {}", lightpath_index, ordered=True)
        lightpath_mask = jnp.where(state.path_index_array == lightpath_index, 0., 1.)  # Allow current lightpath
        #jax.debug.print("lightpath_mask {}", lightpath_mask, ordered=True)
        lightpath_mask = jnp.where(state.path_index_array == -1, 0., lightpath_mask)  # Allow empty slots
        #jax.debug.print("lightpath_mask {}", lightpath_mask, ordered=True)
        lightpath_slots = get_path_slots(lightpath_mask, params, nodes_sd, i)
        #jax.debug.print("lightpath_slots {}", lightpath_slots, ordered=True)
        # Step 3 combine masks
        path_mask = jnp.max(jnp.stack((capacity_slots, lightpath_slots)), axis=0)
        # Swap zeros for ones
        path_mask = jnp.where(path_mask == 0, 1., 0.)
        #jax.debug.print("path_mask {}", path_mask, ordered=True)
        mask = jax.lax.dynamic_update_slice(mask, path_mask, (i * params.link_resources,))
        return mask

    # Loop over each path
    link_slot_mask = jax.lax.fori_loop(0, params.k_paths, mask_path, init_mask)
    if params.aggregate_slots > 1:
        # Full link slot mask is used in process_path_action to get the correct slot from the aggregated slot action
        state = state.replace(full_link_slot_mask=link_slot_mask)
        link_slot_mask, _ = aggregate_slots(link_slot_mask.reshape(params.k_paths, -1), params)
        link_slot_mask = link_slot_mask.reshape(-1)
    state = state.replace(link_slot_mask=link_slot_mask)
    return state

most_used(state, params, unique_lightpaths, relative)

Get the amount of utilised bandwidth on each lightpath. If RWA-LR environment, the utilisation of a slot is defined by either the count of unique active lightpahts on the slot (if unique_lightpaths is True) or the count of active lightpaths on the slot (if unique_lightpaths is False). If RSA-type environment, utilisation is the count of active lightpaths on that slot.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
unique_lightpaths	bool	Whether to consider unique lightpaths	required
relative	bool	Whether to return relative utilisation	required

Returns:

Type	Description
Array	chex.Array: Most used slots (array length = link_resources)

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1, 2, 3))
def most_used(state: EnvState, params: EnvParams, unique_lightpaths, relative) -> chex.Array:
    """Get the amount of utilised bandwidth on each lightpath.
    If RWA-LR environment, the utilisation of a slot is defined by either the count of unique active lightpahts on the
    slot (if unique_lightpaths is True) or the count of active lightpaths on the slot (if unique_lightpaths is False).
    If RSA-type environment, utilisation is the count of active lightpaths on that slot.

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        unique_lightpaths (bool): Whether to consider unique lightpaths
        relative (bool): Whether to return relative utilisation

    Returns:
        chex.Array: Most used slots (array length = link_resources)
    """
    if params.__class__.__name__ != "RWALightpathReuseEnvParams":
        most_used_slots = jnp.count_nonzero(state.link_slot_array, axis=0) + 1
    elif params.__class__.__name__ == "RWALightpathReuseEnvParams" and not unique_lightpaths:
        # Get initial path capacity
        initial_path_capacity = init_path_capacity_array(
            params.link_length_array.val, params.path_link_array.val, scale_factor=1.0
        )
        initial_path_capacity = jnp.squeeze(jax.vmap(lambda x: initial_path_capacity[x])(state.path_index_array))
        utilisation = jnp.where(initial_path_capacity - state.link_capacity_array < 0, 0,
                                initial_path_capacity - state.link_capacity_array)
        if relative:
            utilisation = utilisation / initial_path_capacity
        # Get most used slots by summing the utilisation along the slots
        most_used_slots = jnp.sum(utilisation, axis=0) + 1
    else:
        most_used_slots = jnp.count_nonzero(state.path_index_array + 1, axis=0) + 1
    return most_used_slots

mu_ksp(state, params, unique_lightpaths, relative)

Use the most-used available slot on any path. The most-used slot is that which has the most unique lightpaths (if unique_lightpaths=True) or active lightpaths. Method: Go through action mask and find the usage of available slots, choose available slot that is most utilised.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	EnvParams	Environment parameters	required
unique_lightpaths	bool	Whether to consider unique lightpaths	required
relative	bool	Whether to return relative utilisation	required

Returns:

Type	Description
Array	chex.Array: Action

Source code in xlron/heuristics/heuristics.py

@partial(jax.jit, static_argnums=(1, 2, 3))
def mu_ksp(state: EnvState, params: EnvParams, unique_lightpaths: bool, relative: bool) -> chex.Array:
    """Use the most-used available slot on any path.
    The most-used slot is that which has the most unique lightpaths (if unique_lightpaths=True) or active lightpaths.
    Method: Go through action mask and find the usage of available slots, choose available slot that is most utilised.

    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
        unique_lightpaths (bool): Whether to consider unique lightpaths
        relative (bool): Whether to return relative utilisation

    Returns:
        chex.Array: Action
    """
    mask = get_action_mask(state, params)
    # Get most used slots by summing the link_slot_array along the links
    most_used_slots = most_used(state, params, unique_lightpaths, relative)
    # Get usage of available slots
    most_used_mask = most_used_slots * mask
    # Chosen slot is the most used globally
    action = jnp.argmax(most_used_mask)
    return action

normalise_traffic_matrix(traffic_matrix)

Normalise traffic matrix to sum to 1

Source code in xlron/environments/env_funcs.py

def normalise_traffic_matrix(traffic_matrix):
    """Normalise traffic matrix to sum to 1"""
    traffic_matrix /= jnp.sum(traffic_matrix, promote_integers=False)
    return traffic_matrix

pad_array(array, fill_value)

Pad a ragged multidimensional array to rectangular shape. Used for training on multiple topologies. Source: https://codereview.stackexchange.com/questions/222623/pad-a-ragged-multidimensional-array-to-rectangular-shape

Parameters:

Name	Type	Description	Default
array		array to pad	required
fill_value		value to fill with	required

Returns:

Name	Type	Description
result		padded array

Source code in xlron/environments/env_funcs.py

def pad_array(array, fill_value):
    """
    Pad a ragged multidimensional array to rectangular shape.
    Used for training on multiple topologies.
    Source: https://codereview.stackexchange.com/questions/222623/pad-a-ragged-multidimensional-array-to-rectangular-shape

    Args:
        array: array to pad
        fill_value: value to fill with

    Returns:
        result: padded array
    """

    def get_dimensions(array, level=0):
        yield level, len(array)
        try:
            for row in array:
                yield from get_dimensions(row, level + 1)
        except TypeError: #not an iterable
            pass

    def get_max_shape(array):
        dimensions = defaultdict(int)
        for level, length in get_dimensions(array):
            dimensions[level] = max(dimensions[level], length)
        return [value for _, value in sorted(dimensions.items())]

    def iterate_nested_array(array, index=()):
        try:
            for idx, row in enumerate(array):
                yield from iterate_nested_array(row, (*index, idx))
        except TypeError: # final level
            yield (*index, slice(len(array))), array

    dimensions = get_max_shape(array)
    result = np.full(dimensions, fill_value)
    for index, value in iterate_nested_array(array):
        result[index] = value
    return result

path_action_only(topology_pattern, action_counter, remaining_actions)

This is to check if node has already been assigned, therefore just need to assign slots (n=0)

Parameters:

Name	Type	Description	Default
topology_pattern	Array	Topology pattern	required
action_counter	Array	Action counter	required
remaining_actions	Scalar	Remaining actions	required

Returns:

Name	Type	Description
bool		True if only path action, False if node action

Source code in xlron/environments/env_funcs.py

def path_action_only(topology_pattern: chex.Array, action_counter: chex.Array, remaining_actions: chex.Scalar):
    """This is to check if node has already been assigned, therefore just need to assign slots (n=0)

    Args:
        topology_pattern: Topology pattern
        action_counter: Action counter
        remaining_actions: Remaining actions

    Returns:
        bool: True if only path action, False if node action
    """
    # Get topology segment to be assigned e.g. [2,1,4]
    topology_segment = jax.lax.dynamic_slice(topology_pattern, ((remaining_actions-1)*2, ), (3, ))
    topology_indices = jnp.arange(topology_pattern.shape[0])
    # Check if the latest node in the segment is found in "prev_assigned_topology"
    new_node_to_be_assigned = topology_segment[0]
    prev_assigned_topology = jnp.where(topology_indices > (action_counter[-1]-1)*2, topology_pattern, 0)
    nodes_already_assigned_check = jnp.any(jnp.sum(jnp.where(prev_assigned_topology == new_node_to_be_assigned, 1, 0)) > 0)
    return nodes_already_assigned_check

poisson(key, lam, shape=(), dtype=dtypes.float_)

Sample Exponential random values with given shape and float dtype.

The values are distributed according to the probability density function:

.. math:: f(x) = \lambda e^{-\lambda x}

on the domain :math:0 \le x < \infty.

Args: key: a PRNG key used as the random key. lam: a positive float32 or float64 Tensor indicating the rate parameter shape: optional, a tuple of nonnegative integers representing the result shape. Default (). dtype: optional, a float dtype for the returned values (default float64 if jax_enable_x64 is true, otherwise float32).

Returns: A random array with the specified shape and dtype.

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 2, 3))
def poisson(key: Union[Array, prng.PRNGKeyArray],
            lam: ArrayLike,
            shape: Shape = (),
            dtype: DTypeLike = dtypes.float_) -> Array:
    r"""Sample Exponential random values with given shape and float dtype.

    The values are distributed according to the probability density function:

    .. math::
     f(x) = \lambda e^{-\lambda x}

    on the domain :math:`0 \le x < \infty`.

    Args:
    key: a PRNG key used as the random key.
    lam: a positive float32 or float64 `Tensor` indicating the rate parameter
    shape: optional, a tuple of nonnegative integers representing the result
      shape. Default ().
    dtype: optional, a float dtype for the returned values (default float64 if
      jax_enable_x64 is true, otherwise float32).

    Returns:
    A random array with the specified shape and dtype.
    """
    key, _ = jax._src.random._check_prng_key(key)
    if not dtypes.issubdtype(dtype, np.floating):
        raise ValueError(f"dtype argument to `exponential` must be a float "
                       f"dtype, got {dtype}")
    dtype = dtypes.canonicalize_dtype(dtype)
    shape = core.canonicalize_shape(shape)
    return _poisson(key, lam, shape, dtype)

process_path_action(state, params, path_action)

Process path action to get path index and initial slot index.

Parameters:

Name	Type	Description	Default
state	State	current state	required
params	Params	environment parameters	required
path_action	int	path action	required

Returns:

Name	Type	Description
int	(Array, Array)	path index
int	(Array, Array)	initial slot index

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def process_path_action(state: EnvState, params: EnvParams, path_action: chex.Array) -> (chex.Array, chex.Array):
    """Process path action to get path index and initial slot index.

    Args:
        state (State): current state
        params (Params): environment parameters
        path_action (int): path action

    Returns:
        int: path index
        int: initial slot index
    """
    num_slot_actions = math.ceil(params.link_resources/params.aggregate_slots)
    path_index = jnp.floor(path_action / num_slot_actions).astype(LARGE_INT_DTYPE).reshape(1)
    initial_aggregated_slot_index = jnp.mod(path_action, num_slot_actions).reshape(1)
    initial_slot_index = initial_aggregated_slot_index*params.aggregate_slots
    if params.aggregate_slots > 1:
        # Get the path mask do a dynamic slice and get the index of first unoccupied slot in the slice
        path_mask = jax.lax.dynamic_slice(state.full_link_slot_mask, path_index*params.link_resources, (params.link_resources,))
        path_mask_slice = jax.lax.dynamic_slice(path_mask, initial_slot_index, (params.aggregate_slots,))
        # Use argmax to get index of first 1 in slice of mask
        initial_slot_index = initial_slot_index + jnp.argmax(path_mask_slice).astype(MED_INT_DTYPE)
    return path_index[0], initial_slot_index[0]

read_rsa_request(request_array)

Read RSA request from request array. Return source-destination nodes and bandwidth request.

Parameters:

Name	Type	Description	Default
request_array	Array	request array	required

Returns:

Type	Description
Tuple[Array, Array]	Tuple[chex.Array, chex.Array]: source-destination nodes and bandwidth request

Source code in xlron/environments/env_funcs.py

def read_rsa_request(request_array: chex.Array) -> Tuple[chex.Array, chex.Array]:
    """Read RSA request from request array. Return source-destination nodes and bandwidth request.

    Args:
        request_array: request array

    Returns:
        Tuple[chex.Array, chex.Array]: source-destination nodes and bandwidth request
    """
    node_s = jax.lax.dynamic_slice(request_array, (0,), (1,))
    requested_datarate = jax.lax.dynamic_slice(request_array, (1,), (1,))
    node_d = jax.lax.dynamic_slice(request_array, (2,), (1,))
    nodes_sd = jnp.concatenate((node_s, node_d))
    return nodes_sd, requested_datarate

remove_expired_node_requests(state, params)

Check for values in node_departure_array that are less than the current time but greater than zero (negative time indicates the request is not yet finalised). If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase node_capacity_array by expired resources on each node.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_node_requests(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """Check for values in node_departure_array that are less than the current time but greater than zero
    (negative time indicates the request is not yet finalised).
    If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase
    node_capacity_array by expired resources on each node.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    mask = jnp.where(state.node_departure_array < jnp.squeeze(state.current_time), 1, 0)
    mask = jnp.where(0 < state.node_departure_array, mask, 0)
    expired_resources = jnp.sum(jnp.where(mask == 1, state.node_resource_array, 0), axis=1, promote_integers=False)
    state = state.replace(
        node_capacity_array=state.node_capacity_array + expired_resources,
        node_resource_array=jnp.where(mask == 1, 0, state.node_resource_array),
        node_departure_array=jnp.where(mask == 1, jnp.inf, state.node_departure_array)
    )
    return state

remove_expired_services_rmsa_gn_model(state, params)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rmsa_gn_model(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                      state.link_slot_departure_array - jnp.squeeze(current_time),
                                                      updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=jnp.where(mask_remove == one, zero, state.link_slot_array),
        link_slot_departure_array=updated_link_slot_departure_array,
        link_snr_array=jnp.where(mask_remove == one, zero, state.link_snr_array),
        path_index_array=jnp.where(mask_remove == one, -one, state.path_index_array),
        channel_centre_bw_array=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array),
        channel_power_array=jnp.where(mask_remove == one, zero, state.channel_power_array),
        modulation_format_index_array=jnp.where(mask_remove == one, -one, state.modulation_format_index_array),
        path_index_array_prev=jnp.where(mask_remove == one, -one, state.path_index_array_prev),
        channel_centre_bw_array_prev=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array_prev),
        channel_power_array_prev=jnp.where(mask_remove == one, zero, state.channel_power_array_prev),
        modulation_format_index_array_prev=jnp.where(mask_remove == one, -one, state.modulation_format_index_array_prev),
    )
    return state

remove_expired_services_rsa(state, params)

Check for values in link_slot_departure_array that are less than the current time but greater than zero (negative time indicates the request is not yet finalised). If found, set to zero in link_slot_array and link_slot_departure_array.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rsa(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """Check for values in link_slot_departure_array that are less than the current time but greater than zero
    (negative time indicates the request is not yet finalised).
    If found, set to zero in link_slot_array and link_slot_departure_array.

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_array = jnp.where(mask_remove == one, zero, state.link_slot_array)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                     state.link_slot_departure_array - jnp.squeeze(current_time),
                                                     updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=updated_link_slot_array,
        link_slot_departure_array=updated_link_slot_departure_array,
    )
    return state

remove_expired_services_rsa_gn_model(state, params)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rsa_gn_model(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                      state.link_slot_departure_array - jnp.squeeze(current_time),
                                                      updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=jnp.where(mask_remove == one, zero, state.link_slot_array),
        link_slot_departure_array=updated_link_slot_departure_array,
        link_snr_array=jnp.where(mask_remove == one, zero, state.link_snr_array),
        path_index_array=jnp.where(mask_remove == one, -one, state.path_index_array),
        channel_centre_bw_array=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array),
        channel_power_array=jnp.where(mask_remove == one, zero, state.channel_power_array),
        path_index_array_prev=jnp.where(mask_remove == one, -one, state.path_index_array_prev),
        channel_centre_bw_array_prev=jnp.where(mask_remove == one, zero, state.channel_centre_bw_array_prev),
        channel_power_array_prev=jnp.where(mask_remove == one, zero, state.channel_power_array_prev),
    )
    if params.monitor_active_lightpaths:
        # The active_lightpaths_array is set to -1 when the lightpath is not active
        # The active_lightpaths_array_departure is set to 0 when the lightpath is not active
        # (active_lightpaths_array is used to calculate the total throughput)
        mask_remove = jnp.where(
            (zero <= state.active_lightpaths_array_departure) & (state.active_lightpaths_array_departure <= jnp.squeeze(current_time)),
            one, zero)
        state = state.replace(
            active_lightpaths_array=jnp.where(mask_remove == one, -one, state.active_lightpaths_array),
            active_lightpaths_array_departure=jnp.where(mask_remove == one, zero, state.active_lightpaths_array_departure),
        )
    return state

remove_expired_services_rwalr(state, params)

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required
params	Optional[EnvParams]	Environment parameters	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def remove_expired_services_rwalr(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """

    Args:
        state: Environment state
        params: Environment parameters

    Returns:
        Updated environment state
    """
    # Set one where link_slot_departure_array is >= zero and <= current time
    current_time = state.current_time if not params.relative_arrival_times else state.arrival_time
    mask_remove = jnp.where(
        (zero <= state.link_slot_departure_array) & (state.link_slot_departure_array <= jnp.squeeze(current_time)),
        one, zero)
    updated_link_slot_departure_array = jnp.where(mask_remove == one, zero, state.link_slot_departure_array)
    if params.relative_arrival_times:
        mask_subtract = jnp.where(updated_link_slot_departure_array <= zero, zero, one)
        updated_link_slot_departure_array = jnp.where(mask_subtract == one,
                                                     state.link_slot_departure_array - jnp.squeeze(current_time),
                                                     updated_link_slot_departure_array)
    state = state.replace(
        link_slot_array=jnp.where(mask_remove == one, zero, state.link_slot_array),
        path_index_array=jnp.where(mask_remove == one, -one, state.path_index_array),
        link_slot_departure_array=updated_link_slot_departure_array,
    )
    return state

required_slots(bitrate, se, channel_width, guardband=1)

Calculate required slots for a given bitrate and spectral efficiency.

Parameters:

Name	Type	Description	Default
bit_rate	float	Bit rate in Gbps	required
se	float	Spectral efficiency in bps/Hz	required
channel_width	float	Channel width in GHz	required
guardband	int	Guard band. Defaults to 1.	1

Returns:

Name	Type	Description
int	Array	Required slots

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(2, 3))
def required_slots(bitrate: float, se: int, channel_width: float, guardband: int = 1) -> chex.Array:
    """Calculate required slots for a given bitrate and spectral efficiency.

    Args:
        bit_rate (float): Bit rate in Gbps
        se (float): Spectral efficiency in bps/Hz
        channel_width (float): Channel width in GHz
        guardband (int, optional): Guard band. Defaults to 1.

    Returns:
        int: Required slots
    """
    # If bitrate is zero, then required slots should be zero
    return ((jnp.ceil(bitrate/(se*channel_width))+guardband) * (one - (bitrate == zero))).astype(MED_INT_DTYPE)

set_band_gaps(link_slot_array, params, val)

Set band gaps in link slot array Args: link_slot_array (chex.Array): Link slot array params (RSAGNModelEnvParams): Environment parameters val (int): Value to set Returns: chex.Array: Link slot array with band gaps

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1, 2))
def set_band_gaps(link_slot_array: chex.Array, params: RSAGNModelEnvParams, val: int) -> chex.Array:
    """Set band gaps in link slot array
    Args:
        link_slot_array (chex.Array): Link slot array
        params (RSAGNModelEnvParams): Environment parameters
        val (int): Value to set
    Returns:
        chex.Array: Link slot array with band gaps
    """
    # Create array that is size of link_slot array with values of column index
    mask = jnp.arange(params.link_resources)
    mask = jnp.tile(mask, (params.num_links, 1))
    def set_band_gap(i, arr):
        gap_start = params.gap_starts.val[i]
        gap_end = gap_start + params.gap_widths.val[i]
        condition = jnp.logical_and(arr >= gap_start, arr < gap_end)
        arr = jnp.where(condition, -one, arr)
        return arr
    mask = jax.lax.fori_loop(0, params.gap_widths.val.shape[0], set_band_gap, mask)
    link_slot_array = jnp.where(mask == -one, val, link_slot_array)
    return link_slot_array

undo_action_rmsa_gn_model(state, params)

Undo action for RMSA GN model Args: state (EnvState): Environment state action (chex.Array): Action array params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def undo_action_rmsa_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> EnvState:
    """Undo action for RMSA GN model
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action array
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    state = undo_action_rsa(state, params)  # Undo link_slot_array and link_slot_departure_array
    state = state.replace(
        link_slot_array=set_band_gaps(state.link_slot_array, params, -one),  # Set C+L band gap
        channel_centre_bw_array=state.channel_centre_bw_array_prev,
        path_index_array=state.path_index_array_prev,
        channel_power_array=state.channel_power_array_prev,
        modulation_format_index_array=state.modulation_format_index_array_prev,
    )
    return state

undo_action_rsa(state, params)

If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation. Partial resource allocation is indicated by negative time in link slot departure array. Check for values in link_slot_departure_array that are less than zero. If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def undo_action_rsa(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation.
    Partial resource allocation is indicated by negative time in link slot departure array.
    Check for values in link_slot_departure_array that are less than zero.
    If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    # If departure array is negative, then undo the action
    mask = jnp.where(state.link_slot_departure_array < zero, one, zero)
    # If link slot array is < -1, then undo the action
    # (departure might be positive because existing service had holding time after current)
    # e.g. (time_in_array = t1 - t2) where t2 < t1 and t2 = current_time + holding_time
    # but link_slot_array = -2 due to double allocation, so undo the action
    mask = jnp.where(state.link_slot_array < -one, one, mask)
    state = state.replace(
        link_slot_array=jnp.where(mask == one, state.link_slot_array + one, state.link_slot_array),
        link_slot_departure_array=jnp.where(
            mask == one,
            state.link_slot_departure_array + state.current_time + state.holding_time,
            state.link_slot_departure_array),
        total_bitrate=state.total_bitrate + read_rsa_request(state.request_array)[1][0],
    )
    return state

undo_action_rsa_gn_model(state, params)

Undo action for RSA GN model Args: state (EnvState): Environment state action (chex.Array): Action array params (EnvParams): Environment parameters Returns: EnvState: Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, static_argnums=(1,))
def undo_action_rsa_gn_model(state: RSAGNModelEnvState, params: RSAGNModelEnvParams) -> EnvState:
    """Undo action for RSA GN model
    Args:
        state (EnvState): Environment state
        action (chex.Array): Action array
        params (EnvParams): Environment parameters
    Returns:
        EnvState: Updated environment state
    """
    state = undo_action_rsa(state, params)  # Undo link_slot_array and link_slot_departure_array
    state = state.replace(
        link_slot_array=set_band_gaps(state.link_slot_array, params, -one),  # Set C+L band gap
        channel_centre_bw_array=state.channel_centre_bw_array_prev,
        path_index_array=state.path_index_array_prev,
        channel_power_array=state.channel_power_array_prev,
    )
    if params.monitor_active_lightpaths:
        # If departure array is negative, then undo the action
        mask = jnp.where(state.active_lightpaths_array_departure < zero, one,  zero)
        state = state.replace(
            active_lightpaths_array=jnp.where(mask == one, -one, state.active_lightpaths_array),
            active_lightpaths_array_departure=jnp.where(
                mask == one,
                state.active_lightpaths_array_departure + state.current_time + state.holding_time,
                state.active_lightpaths_array_departure),
        )
    return state

undo_action_rwalr(state, params)

If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation. Partial resource allocation is indicated by negative time in link slot departure array. Check for values in link_slot_departure_array that are less than zero. If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def undo_action_rwalr(state: EnvState, params: Optional[EnvParams]) -> EnvState:
    """If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation.
    Partial resource allocation is indicated by negative time in link slot departure array.
    Check for values in link_slot_departure_array that are less than zero.
    If found, increase link_slot_array by +1 and link_slot_departure_array by current_time + holding_time of current request.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    # If departure array is negative, then undo the action
    mask = jnp.where(state.link_slot_departure_array < zero, one, zero)
    # If link slot array is < -1, then undo the action
    # (departure might be positive because existing service had holding time after current)
    # e.g. (time_in_array = t1 - t2) where t2 < t1 and t2 = current_time + holding_time
    # but link_slot_array = -2 due to double allocation, so undo the action
    mask = jnp.where(state.link_slot_array < -one, one, mask)
    state = state.replace(
        link_slot_array=jnp.where(mask == one, state.link_slot_array + one, state.link_slot_array),
        link_slot_departure_array=jnp.where(
            mask == one,
            state.link_slot_departure_array + state.current_time + state.holding_time,
            state.link_slot_departure_array),
        total_bitrate=state.total_bitrate + read_rsa_request(state.request_array)[1][0]
    )
    return state

undo_node_action(state)

If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation. Partial resource allocation is indicated by negative time in node departure array. Check for values in node_departure_array that are less than zero. If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase node_capacity_array by expired resources on each node.

Parameters:

Name	Type	Description	Default
state	EnvState	Environment state	required

Returns:

Type	Description
EnvState	Updated environment state

Source code in xlron/environments/env_funcs.py

@partial(jax.jit, donate_argnums=(0,))
def undo_node_action(state: EnvState) -> EnvState:
    """If the request is unsuccessful i.e. checks fail, then remove the partial (unfinalised) resource allocation.
    Partial resource allocation is indicated by negative time in node departure array.
    Check for values in node_departure_array that are less than zero.
    If found, set to infinity in node_departure_array, set to zero in node_resource_array, and increase
    node_capacity_array by expired resources on each node.

    Args:
        state: Environment state

    Returns:
        Updated environment state
    """
    # TODO - Check that node resource clash doesn't happen (so time is always negative after implementation)
    #  and undoing always succeeds with negative time
    mask = jnp.where(state.node_departure_array < 0, 1, 0)
    resources = jnp.sum(jnp.where(mask == 1, state.node_resource_array, 0), axis=1, promote_integers=False)
    state = state.replace(
        node_capacity_array=state.node_capacity_array + resources,
        node_resource_array=jnp.where(mask == 1, 0, state.node_resource_array),
        node_departure_array=jnp.where(mask == 1, jnp.inf, state.node_departure_array),
    )
    return state

update_action_history(action_history, action_counter, action)

Update action history by adding action to first available index starting from the end.

Parameters:

Name	Type	Description	Default
action_history	Array	Action history	required
action_counter	Array	Action counter	required
action	Array	Action to add to history	required

Returns:

Type	Description
Array	Updated action_history

Source code in xlron/environments/env_funcs.py

def update_action_history(action_history: chex.Array, action_counter: chex.Array, action: chex.Array) -> chex.Array:
    """Update action history by adding action to first available index starting from the end.

    Args:
        action_history: Action history
        action_counter: Action counter
        action: Action to add to history

    Returns:
        Updated action_history
    """
    return jax.lax.dynamic_update_slice(action_history, jnp.flip(action, axis=0).astype(MED_INT_DTYPE), ((action_counter[-1]-1)*2,))

update_active_lightpaths_array(state, path_index, initial_slot_index, num_slots)

Update active lightpaths array with new path index. Find the first index of the array with value -1 and replace with path index. Args: state (RSAGNModelEnvState): Environment state path_index (int): Path index to add to active lightpaths array Returns: jnp.array: Updated active lightpaths array

Source code in xlron/environments/env_funcs.py

def update_active_lightpaths_array(state: RSAGNModelEnvState, path_index: int, initial_slot_index: int, num_slots: int) -> chex.Array:
    """Update active lightpaths array with new path index.
    Find the first index of the array with value -1 and replace with path index.
    Args:
        state (RSAGNModelEnvState): Environment state
        path_index (int): Path index to add to active lightpaths array
    Returns:
        jnp.array: Updated active lightpaths array
    """
    first_empty_index = jnp.argmin(state.active_lightpaths_array[:, 0])  # Just look at the first column
    return jax.lax.dynamic_update_slice(state.active_lightpaths_array, jnp.array([[path_index, initial_slot_index, num_slots[0]]]), (first_empty_index, 0))

update_active_lightpaths_array_departure(state, time)

Update active lightpaths array with new path index. Find the first index of the array with value -1 and replace with path index. Args: state (RSAGNModelEnvState): Environment state time (float): Departure time Returns: jnp.array: Updated active lightpaths array

Source code in xlron/environments/env_funcs.py

def update_active_lightpaths_array_departure(state: RSAGNModelEnvState, time: float) -> chex.Array:
    """Update active lightpaths array with new path index.
    Find the first index of the array with value -1 and replace with path index.
    Args:
        state (RSAGNModelEnvState): Environment state
        time (float): Departure time
    Returns:
        jnp.array: Updated active lightpaths array
    """
    first_empty_index = jnp.argmin(state.active_lightpaths_array[:, 0])  # Just look at the first column
    return jax.lax.dynamic_update_slice(state.active_lightpaths_array_departure, jnp.stack((time, time, time)), (first_empty_index, 0))

update_graph_tuple(state, params)

Update graph tuple for use with Jraph GNNs. Edge and node features are updated from link_slot_array and node_capacity_array respectively. Global features are updated as request_array. Args: state (EnvState): Environment state params (EnvParams): Environment parameters Returns: state (EnvState): Environment state with updated graph tuple

Source code in xlron/environments/env_funcs.py

def update_graph_tuple(state: EnvState, params: EnvParams):
    """Update graph tuple for use with Jraph GNNs.
    Edge and node features are updated from link_slot_array and node_capacity_array respectively.
    Global features are updated as request_array.
    Args:
        state (EnvState): Environment state
        params (EnvParams): Environment parameters
    Returns:
        state (EnvState): Environment state with updated graph tuple
    """
    # Get source and dest from request array
    source_dest, datarate = read_rsa_request(state.request_array)
    source, dest = source_dest[0], source_dest[2]
    # Global feature is normalised data rate of current request
    globals = jnp.array([datarate / jnp.max(params.values_bw.val)], dtype=LARGE_FLOAT_DTYPE)
    # One-hot encode source and destination
    source_dest_features = jnp.zeros((params.num_nodes, 2), dtype=LARGE_FLOAT_DTYPE)
    source_dest_features = source_dest_features.at[source.astype(MED_INT_DTYPE), 0].set(1)
    source_dest_features = source_dest_features.at[dest.astype(MED_INT_DTYPE), 1].set(-1)
    spectral_features = state.graph.nodes[..., :3]
    holding_time_edge_features = state.link_slot_departure_array / params.mean_service_holding_time

    if params.__class__.__name__ in ["RSAGNModelEnvParams", "RMSAGNModelEnvParams"]:
        # Normalize by max parameters (converted to linear units)
        max_power = isrs_gn_model.from_dbm(params.max_power)
        normalized_power = jnp.round(state.channel_power_array / max_power, 3)
        max_snr = isrs_gn_model.from_db(params.max_snr)
        normalized_snr = jnp.round(state.link_snr_array / max_snr, 3)
        edge_features = jnp.stack([normalized_snr, normalized_power], axis=-1)
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)
    elif params.__class__.__name__ == "VONEEnvParams":
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        node_features = getattr(state, "node_capacity_array", jnp.zeros(params.num_nodes))
        node_features = node_features.reshape(-1, 1)
        node_features = jnp.concatenate([node_features, spectral_features, source_dest_features], axis=-1)
    else:
        edge_features = state.link_slot_array if params.mean_service_holding_time > 1e5 else holding_time_edge_features
        node_features = jnp.concatenate([spectral_features, source_dest_features], axis=-1)

    if params.disable_node_features:
        node_features = jnp.zeros((1,), dtype=LARGE_FLOAT_DTYPE)

    edge_features = edge_features if params.directed_graph else jnp.repeat(edge_features, 2, axis=0)
    graph = state.graph._replace(nodes=node_features, edges=edge_features, globals=globals)
    state = state.replace(graph=graph)
    return state

update_node_array(node_indices, array, node, request)

Used to udated selected_nodes array with new requested resources on each node, for use in

Source code in xlron/environments/env_funcs.py

def update_node_array(node_indices, array, node, request):
    """Used to udated selected_nodes array with new requested resources on each node, for use in """
    return jnp.where(node_indices == node, array-request, array)

vmap_set_path_links(link_slot_array, path, initial_slot, num_slots, value)

Set relevant slots along links in path to val.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
path	Array	Path (row from path-link array that indicates links used by path)	required
initial_slot	int	Initial slot	required
num_slots	int	Number of slots	required
value	int	Value to set on link slot array	required

Returns:

Type	Description
Array	Updated link slot array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_set_path_links(link_slot_array: chex.Array, path: chex.Array, initial_slot: int, num_slots: int, value: int) -> chex.Array:
    """Set relevant slots along links in path to val.

    Args:
        link_slot_array: Link slot array
        path: Path (row from path-link array that indicates links used by path)
        initial_slot: Initial slot
        num_slots: Number of slots
        value: Value to set on link slot array

    Returns:
        Updated link slot array
    """
    return jax.vmap(set_path, in_axes=(0, 0, None, None, None))(link_slot_array, path, initial_slot, num_slots, value)

vmap_update_node_departure(node_departure_array, selected_nodes, value)

Called when implementing node action. Sets request departure time ("value") in place of first "inf" i.e. unoccupied index on node departure array for selected nodes.

Parameters:

Name	Type	Description	Default
node_departure_array	Array	(N x R) Node departure array	required
selected_nodes	Array	(N x 1) Selected nodes (non-zero value on selected node indices)	required
value	int	Value to set on node departure array	required

Returns:

Type	Description
Array	Updated node departure array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_update_node_departure(node_departure_array: chex.Array, selected_nodes: chex.Array, value: int) -> chex.Array:
    """Called when implementing node action.
    Sets request departure time ("value") in place of first "inf" i.e. unoccupied index on node departure array for selected nodes.

    Args:
        node_departure_array: (N x R) Node departure array
        selected_nodes: (N x 1) Selected nodes (non-zero value on selected node indices)
        value: Value to set on node departure array

    Returns:
        Updated node departure array
    """
    first_inf_indices = jnp.argmax(node_departure_array, axis=1).astype(MED_INT_DTYPE)
    return jax.vmap(update_selected_node_departure, in_axes=(0, 0, 0, None))(node_departure_array, selected_nodes, first_inf_indices, value)

vmap_update_node_resources(node_resource_array, selected_nodes)

Called when implementing node action. Sets requested node resources on selected nodes in place of first "zero" i.e. unoccupied index on node resource array for selected nodes.

Parameters:

Name	Type	Description	Default
node_resource_array		(N x R) Node resource array	required
selected_nodes		(N x 1) Requested resources on selected nodes	required

Returns:

Type	Description
	Updated node resource array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_update_node_resources(node_resource_array, selected_nodes):
    """Called when implementing node action.
    Sets requested node resources on selected nodes in place of first "zero" i.e.
    unoccupied index on node resource array for selected nodes.

    Args:
        node_resource_array: (N x R) Node resource array
        selected_nodes: (N x 1) Requested resources on selected nodes

    Returns:
        Updated node resource array
    """
    first_zero_indices = jnp.argmin(node_resource_array, axis=1)
    return jax.vmap(update_selected_node_resources, in_axes=(0, 0, 0))(node_resource_array, selected_nodes, first_zero_indices)

vmap_update_path_links(link_slot_array, path, initial_slot, num_slots, value)

Update relevant slots along links in path to current_val - val.

Parameters:

Name	Type	Description	Default
link_slot_array	Array	Link slot array	required
path	Array	Path (row from path-link array that indicates links used by path)	required
initial_slot	int	Initial slot	required
num_slots	int	Number of slots	required
value	int	Value to subtract from link slot array	required

Returns:

Type	Description
Array	Updated link slot array

Source code in xlron/environments/env_funcs.py

@jax.jit
def vmap_update_path_links(link_slot_array: chex.Array, path: chex.Array, initial_slot: int, num_slots: int, value: int) -> chex.Array:
    """Update relevant slots along links in path to current_val - val.

    Args:
        link_slot_array: Link slot array
        path: Path (row from path-link array that indicates links used by path)
        initial_slot: Initial slot
        num_slots: Number of slots
        value: Value to subtract from link slot array

    Returns:
        Updated link slot array
    """
    return jax.vmap(update_path, in_axes=(0, 0, None, None, None))(link_slot_array, path, initial_slot, num_slots, value)