Project: s41467-024-49261-6 | Function: Heun_step (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 Heun predictor-corrector method for numerically solving the Landau-Lifshitz-Gilbert equation. This second-order numerical scheme evolves the magnetization dynamics with improved accuracy compared to Euler method.

**KEY TERMS:** {
  "M0": "current magnetization state M₀",
  "Mp": "predictor step magnetization M_p",
  "Mc": "corrector step magnetization M_c",
  "dt": "time step Δt",
  "dM": "LLG equation right-hand side dM/dt"
}

**CONTEXT CODE:** 
#
# Hamiltonian class 
#
class Hamiltonian:
	def __init__(self, theta, J0, K0, K1, *, relaxed = 1.0, J_perp_scale=1.0):
		self.theta = theta
		theta_rad = 2*math.pi/360*theta # tilted angle in radian
		if theta == 0:
			self.N0 = 50
		else:
			self.N0 = int(1/theta_rad) # approximate Moire periodicity
		self.J0 = J0
		self.K0 = K0
		self.K1 = K1
		self.Bext = np.zeros((3,))
		self.J_perp_scale = J_perp_scale
		self.relaxed = relaxed

		J_perp_rigid = set_inter_layer_coupling('deltaE_Nano_2.csv', self.N0)
		J_perp_relax = set_inter_layer_coupling('deltaE_Nano_3.csv', self.N0)
		self.J_perp = (J_perp_rigid*(1-relaxed) + J_perp_relax*relaxed)*self.J_perp_scale
		if theta == 0:
			self.J_perp = self.J_perp[0,0,0,0]
		self.total_sites = 4*self.N0*self.N0*2

	def update_B_field(self, B_new):
		self.Bext = B_new

	#
	# get the effective magnetic fields
	#
	def get_Heff(self, M0):
		"""Computes the effective magnetic field in the anisotropic Heisenberg model for twisted double-bilayer CrI3."""
		raise NotImplementedError("Function has been masked for testing purposes")


	#
	# the right-hand side of the LLG equation, M-dot
	#
	def dM(self, M0, Zs):
		Z1, Z2 = Zs
		tau = self.get_torque(M0)
		g0 = self.get_MMH(M0)
		dM = -Z1*tau - Z2*g0
		return dM

	#
	# time-evolution by the Heun scheme 
	#
	def Heun_step(self, M0, dt, Zs):
		"""Implements the Heun predictor-corrector method for numerically solving the Landau-Lifshitz-Gilbert equation."""
		raise NotImplementedError("Function has been masked for testing purposes")

	#
	# time-evolution by the RK45
	#
	def RK45_step(self, M0, dt, Zs):
		k1 = self.dM(M0, Zs)
		k2 = self.dM(M0 + 0.5*dt*k1, Zs)
		k3 = self.dM(M0 + 0.5*dt*k2, Zs)
		k4 = self.dM(M0 + dt*k3, Zs)
		
		M1 = M0 + 1.0/6.0*dt*(k1+2*k2+2*k3+k4)

		return M1



	#
	# cross product of M and H
	#
	def get_torque(self, M0):
		H0 = self.get_Heff(M0)
		tau = np.empty_like(M0)

		tau[...,0] = M0[...,1]*H0[...,2]- M0[...,2]*H0[...,1]
		tau[...,1] = M0[...,2]*H0[...,0]- M0[...,0]*H0[...,2]
		tau[...,2] = M0[...,0]*H0[...,1]- M0[...,1]*H0[...,0]

		return tau


	#
	# get the steepest descent with conserved norm: Mx(MxH) = (M.H)M-M^2H, Eq.(3.25), (3.26) in Exl's thesis
	#
	def get_MMH(self, M0):
		H0 = self.get_Heff(M0)
		g0 = np.empty_like(H0)
				
		# muS_list = np.linalg.norm(M0[:,0,0,0,:], axis =1)
		# for ii in range(4):
		# 	g0[ii,...] = -H0[ii,...]*muS_list[ii]**2

		muS_squared = np.linalg.norm(M0, axis = -1, keepdims=True)**2
		g0 = -muS_squared*H0 # array broadcasting

		MH = (M0[...,2]*H0[...,2] + M0[...,1]*H0[...,1]+ M0[...,0]*H0[...,0]) 
		g0[...,0] +=  M0[...,0]*MH
		g0[...,1] +=  M0[...,1]*MH
		g0[...,2] +=  M0[...,2]*MH

		return g0



	#
	# get energy
	#
	def get_energy(self, M0):

		H0 = self.get_Heff(M0)
		energy = 0.0
		N3, N1, N2, _, _ = np.shape(M0)

		energy = -np.sum(H0*M0)/2

		Zeemann = np.sum(M0*self.Bext) # broadcasting the arrays
		
		energy = energy - Zeemann/2
		energy = energy/self.total_sites

		return energy




	#
	# semi-implicite method with fixed B
	#
	def SIM_minimization_steps(self, S0, Nt_min, verbose=False):
		tol = 10.0**(-8)

		# Note that changing M0 now changes S0.M
		M0 = S0.M
		
		dt = 1/np.sum(M0*self.get_Heff(M0))
		M1 =self.SIM_update_step(M0, dt)
		
		tt = 0
		g0 = 1.0
		last_energy = 10.0*np.ones(20)
		e1 = self.get_energy(M1)

		while g0 > tol or tt < Nt_min:

			dt = self.get_backtracking_time_steps(M0, 0.000001, 0.7)

			M0 = M1
			M1 = self.SIM_update_step(M0,dt)


Complete the function implementation replacing "raise NotImplementedError".