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
10import jax.numpy as jnp
11import jax
12# import matplotlib.pyplot as plt
13
14
15import autohf
16
17
18hops = (np.array([ 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2,
19 2, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5,
20 5, 5, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 8, 8, 8,
21 8, 8, 8, 9, 9, 9, 9, 9, 9, 10, 10, 10, 10, 10, 10, 11, 11,
22 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 14,
23 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16, 16, 16,
24 17, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19,
25 19, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, 21, 21, 22, 22, 22, 22,
26 22, 22, 23, 23, 23, 23, 23, 23, 24, 24, 24, 24, 24, 24, 25, 25, 25,
27 25, 25, 25, 26, 26, 26, 26, 26, 26, 27, 27, 27, 27, 27, 27, 28, 28,
28 28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30, 30, 31,
29 31, 31, 31, 31, 31, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33, 33,
30 34, 34, 34, 34, 34, 34, 35, 35, 35, 35, 35, 35]),
31 np.array([ 1, 5, 6, 11, 30, 31, 0, 2, 6, 7, 31, 32, 1, 3, 7, 8, 32,
32 33, 2, 4, 8, 9, 33, 34, 3, 5, 9, 10, 34, 35, 0, 4, 10, 11,
33 30, 35, 0, 1, 7, 11, 12, 17, 1, 2, 6, 8, 12, 13, 2, 3, 7,
34 9, 13, 14, 3, 4, 8, 10, 14, 15, 4, 5, 9, 11, 15, 16, 0, 5,
35 6, 10, 16, 17, 6, 7, 13, 17, 18, 23, 7, 8, 12, 14, 18, 19, 8,
36 9, 13, 15, 19, 20, 9, 10, 14, 16, 20, 21, 10, 11, 15, 17, 21, 22,
37 6, 11, 12, 16, 22, 23, 12, 13, 19, 23, 24, 29, 13, 14, 18, 20, 24,
38 25, 14, 15, 19, 21, 25, 26, 15, 16, 20, 22, 26, 27, 16, 17, 21, 23,
39 27, 28, 12, 17, 18, 22, 28, 29, 18, 19, 25, 29, 30, 35, 19, 20, 24,
40 26, 30, 31, 20, 21, 25, 27, 31, 32, 21, 22, 26, 28, 32, 33, 22, 23,
41 27, 29, 33, 34, 18, 23, 24, 28, 34, 35, 0, 5, 24, 25, 31, 35, 0,
42 1, 25, 26, 30, 32, 1, 2, 26, 27, 31, 33, 2, 3, 27, 28, 32, 34,
43 3, 4, 28, 29, 33, 35, 4, 5, 24, 29, 30, 34])
44 ) # fmt: skip
45
46
47#### Define Parameters
48Nx, Ny = 6, 6
49U = 4
50Ne = 18
51steps = 2000
52batch_size = 4
53
54hf_settings = {
55 "verbose": False,
56 "steps": steps,
57 "opt_method": "lbfgs",
58 "ansatz": "SD_ROT",
59 "batch_size": batch_size,
60 "gpu": False,
61 "nelec": (Ne, Ne),
62 "noncollinear": True,
63 "state0_scale": 0.01,
64 "seed": 1234,
65}
66
67
68# Build a lagrange multiplier for total Sz
69# We'll use this so that our energy is E+lambda*|\sum_i Sz_i|^2+100
70# Importante note: as an example, let's shift the energy
71# so we understand when its the special loss vs real energy
72lam = 10.0
73
74
75def Esz2(state, rdms):
76 r_uu, r_dd = rdms[0], rdms[1]
77 return lam * jnp.abs(jnp.sum(r_uu.diagonal() - r_dd.diagonal())) ** 2 + 100
78
79
80##
81## We'll do the following steps:
82## 1. run a calculation with no steps to get an energy function
83## because our energy function will really be a loss,
84## polluted with the lagrange multiplier
85##
86## 2. For illustration, we'll run without the multiplier,
87## to see how the the GHF state can rotate arbitrarily
88## You may naturally do this step in order to realize
89## the state needs constraining more explicitly
90##
91## 3. We run AutoHf with the custom energy term
92##
93## 4. Compare 2 and 3
94##
95
96#### Setup calculation
97N = Nx * Ny
98assert Nx == 6 and Ny == 6 # harded coded here
99T = np.zeros((N, N))
100T[hops] = -1
101
102H = autohf.AutoHFHamiltonian(T=(T, T), U=U)
103
104####### Step 1 ######
105# First get energy function
106temp_settings = hf_settings | {"steps": -1}
107data, data_functions = autohf.solve_hf(
108 H,
109 settings=temp_settings,
110 verbosity=0, # quiet mode
111)
112energy_func = data_functions["energy_func"]
113
114####### Step 2 ######
115# run with default settings, no extra term
116data, data_functions = autohf.solve_hf(
117 H,
118 settings=hf_settings,
119)
120
121final_rdm = data_functions["makeRDMs"](data["state"])
122r_uu, r_dd = final_rdm[:N, :N], final_rdm[N:, N:]
123
124####### Step 3 ######
125# run with addition term
126# for illustration, energy is shifted by 100
127data_sz, data_functions_sz = autohf.solve_hf(
128 H,
129 settings=hf_settings,
130 custom_energy_term=Esz2,
131)
132
133final_rdm_sz = data_functions_sz["makeRDMs"](data_sz["state"])
134r_uu_sz, r_dd_sz = final_rdm_sz[:N, :N], final_rdm_sz[N:, N:]
135
136####### Step 4 ######
137print("True energy: ", energy_func(data["state"]))
138print("True energy w/ term:", energy_func(data_sz["state"]))
139print("Final sz: ", np.sum(r_uu.diagonal() - r_dd.diagonal()))
140print("Final sz w/ term:", np.sum(r_uu_sz.diagonal() - r_dd_sz.diagonal()))