Project: s41467-025-60878-z | Function: SEAdLIF.step_fn (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:** Implements the Symplectic Euler discretized adaptive LIF neuron step function. This function performs the core neuron dynamics update including membrane potential, adaptation current, and spike generation according to the SE-adLIF equations from the paper.

**KEY TERMS:** {
  "u": "membrane potential",
  "w": "adaptation current",
  "alpha": "membrane time constant decay",
  "beta": "adaptation time constant decay",
  "a": "sub-threshold adaptation parameter",
  "b": "spike-triggered adaptation parameter",
  "z": "spike output"
}

**CONTEXT CODE:**         self.alpha = cfg.get('alpha', 5.0)
        self.c = cfg.get('c', 0.4)
        self.tau_u_range = cfg.tau_u_range
        self.train_tau_u_method = cfg.get("train_tau", 'interpolation')
        self.tau_w_range = cfg.tau_w_range
        self.train_tau_w_method = cfg.get("train_tau", 'interpolation')        
        self.use_recurrent = cfg.get('use_recurrent', True)
        
        self.ff_gain = cfg.get('ff_gain', 1.0)
        self.a_range =  cfg.get('a_range', [0.0, 1.0])
        self.b_range = cfg.get('b_range',[0.0, 2.0])
        self.num_out_neuron = cfg.get('num_out_neuron', self.out_features)
        self.use_u_rest = cfg.get('use_u_rest', False)
        self.train_u0 = cfg.get('train_u0', False)

        self.q = cfg.q
        
        self.tau_u_trainer: TauTrainer = get_tau_trainer_class(self.train_tau_u_method)(
                self.out_features,
                self.dt, 
                self.tau_u_range[0], 
                self.tau_u_range[1],
                **factory_kwargs)
        
        self.tau_w_trainer: TauTrainer = get_tau_trainer_class(self.train_tau_w_method)(
                self.out_features,
                self.dt, 
                self.tau_w_range[0], 
                self.tau_w_range[1],
                **factory_kwargs)
        
        
        self.weight = Parameter(
            torch.empty((self.out_features, self.in_features), **factory_kwargs)
        )
        self.bias = Parameter(torch.empty(self.out_features, **factory_kwargs))
        
        if self.use_recurrent:
            self.recurrent = Parameter(
                torch.empty((self.out_features, self.out_features), **factory_kwargs)
            )
        else:
            # registering an empty size tensor is required for the static analyser when using jit.script
            self.register_buffer("recurrent", torch.empty(size=()))

        self.a = Parameter(torch.empty(self.out_features, **factory_kwargs))
        self.b = Parameter(torch.empty(self.out_features, **factory_kwargs))
        self.u0 = Parameter(torch.empty(self.out_features, **factory_kwargs), requires_grad=self.train_u0)

        self.reset_parameters()
        def step_fn(recurrent, alpha, beta, thr, a, b, u_rest, carry, cur):
            """
            Implements the Symplectic Euler discretized adaptive LIF neuron step function.
            This function performs the core neuron dynamics update including membrane potential,
            adaptation current, and spike generation according to the SE-adLIF equations from the paper.
            """
            # MASKED: Implements the Symplectic Euler discretized adaptive LIF neuron step function with membrane potential update, adaptation current computation, and spike generation according to SE-adLIF equations for spiking neural network dynamics implementation
            raise NotImplementedError("This function is masked for testing purposes.")
        self.step = step_fn
        
        def wrapped_scan(u0: Parameter, z0: Tensor, w0: Parameter, 
                         x: Tensor,
                         recurrent: Parameter, alpha: Parameter, beta: Parameter, 
                         thr: Tensor, a: Parameter, b: Parameter):
            if self.use_u_rest:
                u_rest = u0
            else:
                u_rest = torch.zeros_like(u0)
            def wrapped_step(carry, cur):
                return step_fn(recurrent, alpha, beta, thr, a, b, u_rest, carry, cur)

            return generic_scan(wrapped_step, (u0, z0, w0), x, self.unroll)
        def wrapped_scan_with_states(u0: Parameter, z0: Tensor, w0: Parameter, x: Tensor,
                         recurrent: Parameter, alpha: Parameter, beta: Parameter, 
                         thr: Tensor, a: Parameter, b: Parameter):
            if self.use_u_rest:
                u_rest = u0
            else:
                u_rest = torch.zeros_like(u0)
            def wrapped_step(carry, cur):
                return step_fn(recurrent, alpha, beta, thr, a, b, u_rest, carry, cur)

            return generic_scan_with_states(wrapped_step, (u0, z0, w0), x, self.unroll)
        self.wrapped_scan = wrapped_scan
        self.wrapped_scan_with_states = wrapped_scan_with_states
        if not cfg.get('compile', False):
            self.wrapped_scan = torch.compiler.disable(self.wrapped_scan, recursive=False)
            self.wrapped_scan_with_states = torch.compiler.disable(self.wrapped_scan_with_states, recursive=False)
        self.reset_parameters()
        
    def reset_parameters(self) -> None:
        self.tau_u_trainer.reset_parameters()
        self.tau_w_trainer.reset_parameters()
        
        
        torch.nn.init.uniform_(
            self.weight,
            -self.ff_gain * torch.sqrt(1 / torch.tensor(self.in_features)),
            self.ff_gain * torch.sqrt(1 / torch.tensor(self.in_features)),
        )
        
        torch.nn.init.zeros_(self.bias)
        
        # h0 states 
        if self.train_u0:
            torch.nn.init.uniform_(self.u0, 0, self.thr[0].item())
        else:
            torch.nn.init.zeros_(self.u0)
        if self.use_recurrent:
            torch.nn.init.orthogonal_(
                self.recurrent,
                gain=1.0,
            )
        
        torch.nn.init.uniform_(self.a, self.a_range[0], self.a_range[1])
        torch.nn.init.uniform_(self.b, self.b_range[0], self.b_range[1])
        
    def initial_state(self, batch_size:int, device: Optional[torch.device] = None) -> tuple[Tensor, Tensor, Tensor]:
        size = (batch_size, self.out_features)
        z = torch.zeros(
            size=size, 
            device=device, 
            dtype=torch.float, 
            layout=None, 
            pin_memory=None,
            requires_grad=True
        )
        w = torch.zeros(
            size=size,
            device=device, 
            dtype=torch.float, 
            layout=None, 
            pin_memory=None,
            requires_grad=True
        )
        return self.u0.unsqueeze(0), z, w

    def apply_parameter_constraints(self):
        self.tau_u_trainer.apply_parameter_constraints()
        self.tau_w_trainer.apply_parameter_constraints()
        self.a.data = torch.clamp(self.a, min=self.a_range[0], max=self.a_range[1])
        self.b.data = torch.clamp(self.b, min=self.b_range[0], max=self.b_range[1])
        self.u0.data = self.u0 - torch.sign(self.u0)*torch.relu(torch.abs(self.u0) - self.thr)
        self.thr.data = torch.maximum(self.thr, torch.zeros_like(self.thr))
        
    def forward(
        self, input_tensor: Tensor,  states: Tuple[Tensor, Tensor, Tensor]
    ) -> Tuple[Tensor, Tensor]:
        decay_u = self.tau_u_trainer.get_decay()
        decay_w = self.tau_w_trainer.get_decay()


Complete the function implementation replacing "raise NotImplementedError".