Project: s41467-024-49261-6 | Function: get_Heff (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:** Computes the effective magnetic field in the anisotropic Heisenberg model for twisted double-bilayer CrI3. Calculates intralayer exchange coupling, interlayer coupling with spatially varying moiré pattern, magnetic anisotropy terms, and external field contributions.

**KEY TERMS:** {
  "J0": "intralayer exchange coupling J₀",
  "J_perp": "spatially varying interlayer coupling J⊥(r)",
  "K0": "exchange anisotropy parameter K₀",
  "K1": "uniaxial anisotropy parameter K₁",
  "Bext": "external magnetic field B_ext",
  "spin_a": "sublattice A spins",
  "spin_b": "sublattice B spins"
}

**CONTEXT CODE:** from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import math
import matplotlib.transforms as mtransforms

muS = 1.5 # magnetization on Cr atoms (muB)
a0 = 6.83 # lattice consant of a hexagonal lattice (angstrom) 
d0 = 6.83/np.sqrt(3) # distance between neighboring sites of a hexagonal lattice (angstrom) 

a1 = a0*np.array([-1/2, np.sqrt(3)/2, 0]) # Bravais lattice vector 
a2 = a0*np.array([1/2, np.sqrt(3)/2, 0]) # Bravais lattice vector
a3 = np.array([0,0, 6.61]) # Bravais lattice vector


def get_atom_position(l1, l2, l3, b):
	return l1*a1+l2*a2+l3*a3+b*[0,d0,0]


#
# 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)


Complete the function implementation replacing "raise NotImplementedError".