Project: s41467-024-44900-4 | Function: set_activation_function (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Sets the nonlinear activation function for reservoir units including tanh, sigmoid, relu and other neural activation functions

**KEY TERMS:** {
  "tanh": "hyperbolic tangent activation",
  "sigmoid": "sigmoid activation function",
  "relu": "rectified linear unit",
  "leaky_relu": "leaky rectified linear unit"
}

**CONTEXT CODE:**         """

        super().__init__(*args, **kwargs)

        # activation function
        self.activation_function = self.set_activation_function(
            activation_function)
        
        # leak rate
        self.leak_rate = leak_rate

    def simulate(
        self, ext_input, w_in, ic=None, output_nodes=None,
        return_states=True, **kwargs
    ):
        """
        Simulates reservoir dynamics given an external input signal
        'ext_input' and an input connectivity matrix 'w_in'

        Parameters
        ----------
        ext_input : (time, N_inputs) numpy.ndarray
            External input signal
            N_inputs: number of external input signals
        w_in : (N_inputs, N) numpy.ndarray
            Input connectivity matrix (source, target)
            N_inputs: number of external input signals
            N: number of nodes in the network
        ic : (N,) numpy.ndarray, optional
            Initial conditions
            N: number of nodes in the network. If w is directed, then rows
            (columns) should correspond to source (target) nodes.
        output_nodes : list or numpy.ndarray, optional
            List of nodes for which reservoir states will be returned if
            'return_states' is True.
        return_states : bool, optional
            If True, simulated resrvoir states are returned. True by default.
        kwargs:
            Other keyword arguments are passed to self.activation_function

        Returns
        -------
        self._state : (time, N) numpy.ndarray
            Activation states of the reservoir.
            N: number of nodes in the network if output_nodes is None, else
            number of output_nodes
        """
        # MASKED: Core reservoir computing simulation that propagates external input through recurrent neural network dynamics with nonlinear activation functions implementation
        raise NotImplementedError("This function is masked for testing purposes.")

    def set_activation_function(self, function):
        # MASKED: Sets the nonlinear activation function for reservoir units including tanh, sigmoid, relu and other neural activation functions implementation
        raise NotImplementedError("This function is masked for testing purposes.")


class SpikingNeuralNetwork(Reservoir):
    """
    Class that represents a Spiking Neural Network
    Adapted from Kim et al., 2019 and Nicola & Clopath, 2017
    (https://github.com/rkim35/spikeRNN/blob/master/spiking/LIF_network_fnc.m)
    ...

    Attributes
    ----------
    w : (N, N) numpy.ndarray
        reservoir connectivity matrix (source, target)
        N: number of nodes in the network. If w is directed, then rows
        (columns) should correspond to source (target) nodes.
    _state : same as ext_input
        reservoir activation states
    n_nodes : int
        dimension of the reservoir
    inh : (N,) numpy.ndarray
        boolean array indicating whether a node is
        inhibitory (True) or excitatory (False)
        N: number of nodes in the network
    exc : (N,) numpy.ndarray
        boolean array indicating whether a node is
        excitatory (True) or inhibitory (False)
        N: number of nodes in the network
    som : (N,) numpy.ndarray
        boolean array indicating whether a node is
        somatostatin-expressing (True) or not (False)
        N: number of nodes in the network
    dt : float
        sampling rate (in s)
    T : float
        trial duration (in s)
    nt : int
        number of time steps
    td : float or numpy.ndarray
        decay time constants of the synaptic filter model (in s)
    REC : (nt, N) numpy.ndarray
        membrane voltage tracings (mV)
        nt: number of time steps
        N: number of nodes in the network
    Is : (N, nt) numpy.ndarray
        external input current
        nt: number of time steps
        N: number of nodes in the network
    IPSCs : (N, nt) numpy.ndarray
        post synaptic currents over time
        nt: number of time steps
        N: number of nodes in the network
    spk : (N, nt) numpy.ndarray
        spike raster
        nt: number of time steps
        N: number of nodes in the network
    rs : (N, nt/timescale) numpy.ndarray
        filtered firing rates over time
        nt: number of time steps
        N: number of nodes in the network
        timescale: number of internal time steps per external time steps
    hs : (N, nt) numpy.ndarray
        filtered firing rates over time (synaptic input accumulation)
        nt: number of time steps
        N: number of nodes in the network
    tspike : (N_spikes, 2) numpy.ndarray
        spike times (in s)
        N_spikes: number of spikes
        tspike[:, 0]: spike neuronal indices
        tspike[:, 1]: spike times (in s)
    inh_fr : (N_inh,) numpy.ndarray
        average firing rates of inhibitory neurons
        N_inh: number of inhibitory neurons
    exc_fr : (N_exc,) numpy.ndarray
        average firing rates of excitatory neurons
        N_exc: number of excitatory neurons
    all_fr : (N,) numpy.ndarray
        average firing rates of all neurons
        N: number of neurons in the network

    Methods
    -------
    # TODO

    simulate

    """

    def __init__(self, *args, inh = 0.2, som = 0., apply_Dale = True, **kwargs):
        """
        Constructor class for Spiking Neural Networks

        Parameters
        ----------
        w: (N, N) numpy.ndarray
            Reservoir connectivity matrix (source, target)
            N: number of nodes in the network. If w is directed,
            then rows (columns) should correspond to source (target) nodes.


Complete the function implementation replacing "raise NotImplementedError".