1import os
2
3os.environ["JAX_PLATFORMS"] = "cpu"
4os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
5from jax import config
6
7config.update("jax_enable_x64", True)
8
9import numpy as np
10# import matplotlib.pyplot as plt
11
12
13import autohf
14
15
16def makeT(Nx, Ny, xperiodic=True, yperiodic=True):
17 N = Nx * Ny
18 T = np.zeros([N, N])
19 # fill hopping matrix
20 for i in range(Nx):
21 for j in range(Ny):
22 a = i * Ny + j
23 if i + 1 < Nx or yperiodic:
24 b = ((i + 1) % Nx) * Ny + j
25
26 T[a, b] = -1
27 T[b, a] = -1
28 if i + 1 < Ny or xperiodic:
29 b = i * Ny + (j + 1) % Ny
30 T[a, b] = -1
31 T[b, a] = -1
32 return T
33
34
35#### Define Parameters
36Nx, Ny = 4, 4
37M = 0.15
38U = 1.5
39Ne = 7
40steps = 100
41batch_size = 4
42
43hf_settings = {
44 "verbose": False,
45 "steps": steps,
46 "opt_method": "lbfgs",
47 "ansatz": "SD_ROT",
48 "batch_size": batch_size,
49 "gpu": False,
50 "nelec": (Ne, Ne),
51 "noncollinear": True,
52}
53
54#### Setup calculation
55N = Nx * Ny
56T = makeT(Nx, Ny)
57
58# Convert to GHF form
59T_ghf = np.array(autohf.tools.uhf_to_ghf([T, T])) # returns jax
60# add M cdag_up c_dn
61for i in range(N):
62 T_ghf[i, i + N] = M
63 T_ghf[i + N, i] = M
64
65# check legal
66assert np.allclose(T_ghf, T_ghf.conj().T)
67
68# Because this is GHF we need to pass L so that we get the correct
69# N can also be passed, which would be N=T_ghf.shape[0]//2
70#
71# For GHF we either pass (T_ghf,) or a (1,2N,2N) array
72H = autohf.AutoHFHamiltonian(settings={"L": (Nx, Ny)}, T=(T_ghf,), U=U)
73
74autohf.solve_hf(
75 H,
76 settings=hf_settings,
77)