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
37U = 1.5
38Ne = 7
39steps = 100
40batch_size = 4
41
42hf_settings = {
43 "verbose": False,
44 "steps": steps,
45 "opt_method": "lbfgs",
46 "ansatz": "SD_ROT",
47 "batch_size": batch_size,
48 "gpu": False,
49 "nelec": (6, 2),
50 "noncollinear": False,
51}
52
53#### Setup calculation
54N = Nx * Ny
55T = makeT(Nx, Ny)
56
57H = autohf.AutoHFHamiltonian(T=(T, T), U=U)
58
59data, data_F = autohf.solve_hf(
60 H,
61 settings=hf_settings,
62)
63
64data, data_F = autohf.solve_hf(
65 H, settings=hf_settings | {"ansatz": "SD", "riemannian": True}, restart_from=(data, data_F)
66)
67
68print([x.shape for x in data["orbitals"]])
69rdms = data_F["makeRDMs"](data["state"])
70print("<Sz> =", np.round((rdms[0].diagonal() - rdms[1].diagonal()).sum() / 2, 4))
71print("<n> =", np.round((rdms[0].diagonal() + rdms[1].diagonal()).sum(), 4))