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 jax
10import jax.numpy as jnp
11
12import numpy as np
13# import matplotlib.pyplot as plt
14
15
16import autohf
17
18
19def makeT(Nx, Ny, xperiodic=True, yperiodic=True):
20 N = Nx * Ny
21 T = np.zeros([N, N])
22 # fill hopping matrix
23 for i in range(Nx):
24 for j in range(Ny):
25 a = i * Ny + j
26 if i + 1 < Nx or yperiodic:
27 b = ((i + 1) % Nx) * Ny + j
28
29 T[a, b] = -1
30 T[b, a] = -1
31 if i + 1 < Ny or xperiodic:
32 b = i * Ny + (j + 1) % Ny
33 T[a, b] = -1
34 T[b, a] = -1
35 return T
36
37
38Nx, Ny = 4, 4
39N = Nx * Ny
40Ne = 8
41T = makeT(Nx, Ny)
42
43H = autohf.AutoHFHamiltonian(T=(T, T), U=4)
44
45TGHF_ = autohf.tools.uhf_to_ghf([T, T])
46signs = (-1) ** (np.arange(N) // Nx + np.arange(N) % Nx)
47
48
49def orbitalFunc(state):
50 # take the hopping matrix TGHF, and "set" M on the off-diagonal blocks
51 TGHF = TGHF_ + 0.0
52 TGHF = TGHF.at[N + np.arange(N), np.arange(N)].set(signs * state)
53 TGHF = TGHF.at[np.arange(N), N + np.arange(N)].set(signs * state)
54 lams, vecs = jnp.linalg.eigh(TGHF)
55 return jnp.stack([vecs[:, : 2 * Ne]])
56
57
58def rdmFunc(state):
59 o = orbitalFunc(state)
60 return autohf.solver.sd_to_rdm_noncollinear(o)
61
62
63hf_settings = {
64 "verbose": False,
65 "steps": 2000,
66 "opt_method": "lbfgs",
67 "ansatz": "CUSTOM",
68 "batch_size": 16,
69 "gpu": False,
70 "nelec": (Ne, Ne),
71 "noncollinear": True,
72}
73
74rng = np.random.default_rng(42)
75state0_ref = 0.0 * np.arange(N).reshape(N)
76state0 = state0_ref + rng.normal(scale=0.01, size=(hf_settings["batch_size"], N))
77
78autohf.solve_hf(
79 H,
80 settings=hf_settings,
81 state2orbitals=orbitalFunc,
82 state2rdm=rdmFunc,
83 state0=state0,
84 state0_ref=state0_ref,
85 jaxoptargs=dict(tol=1e-14), # secret sauce
86)