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
12import autohf
13import jax
14import optax
15from typing import NamedTuple, Union
16import optax.tree_utils as otu
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 = 8
50Ne = 18
51steps = 500
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 "seed": 42,
62 "nelec": (Ne, Ne),
63 "noncollinear": True,
64 "state0_scale": 0.001, # start close to U=0
65}
66
67#### Setup calculation
68N = Nx * Ny
69assert Nx == 6 and Ny == 6 # harded coded here
70T = np.zeros((N, N))
71T[hops] = -1
72
73H = autohf.AutoHFHamiltonian(T=(T, T), U=U)
74
75## Standard lbfgs to get a sense of things
76data, data_f = autohf.solve_hf(
77 H,
78 settings=hf_settings,
79)
80E_lbfgs = data["E"]
81
82# Gradient Descent via optax
83opt = optax.chain(
84 # Sets the parameters of Adam. Note the learning_rate is not here.
85 optax.scale_by_adam(b1=0.9, b2=0.999, eps=1e-8),
86 # Puts a minus sign to *minimize* the loss.
87 optax.scale(-0.01),
88)
89
90data, data_f = autohf.solve_hf(
91 H,
92 settings=hf_settings | dict(opt_method="grad"),
93 optimizer=opt,
94)
95E_grad = data["E"]
96
97
98# Now Construct custom optimization ala Optax examples
99def run_opt(init_params, fun, opt, max_iter, tol):
100 value_and_grad_fun = optax.value_and_grad_from_state(fun)
101
102 def step(carry):
103 params, state = carry
104 value, grad = value_and_grad_fun(params, state=state)
105 updates, state = opt.update(grad, state, params, value=value, grad=grad, value_fn=fun)
106 params = optax.apply_updates(params, updates)
107 return params, state
108
109 def continuing_criterion(carry):
110 _, state = carry
111 iter_num = otu.tree_get(state, "count")
112 grad = otu.tree_get(state, "grad")
113 err = otu.tree_l1_norm(grad)
114 return (iter_num == 0) | ((iter_num < max_iter) & (err >= tol))
115
116 init_carry = (init_params, opt.init(init_params))
117 final_params, final_state = jax.lax.while_loop(continuing_criterion, step, init_carry)
118 return final_params, final_state
119
120
121class InfoState(NamedTuple):
122 iter_num: Union[np.number, np.ndarray, jax.Array]
123
124
125def print_info():
126 def init_fn(params):
127 del params
128 return InfoState(iter_num=0)
129
130 def update_fn(updates, state, params, *, value, grad, **extra_args):
131 del params, extra_args
132
133 jax.debug.print(
134 "Iteration: {i}, Value: {v}, Gradient norm: {e}",
135 i=state.iter_num,
136 v=value,
137 e=otu.tree_l1_norm(grad),
138 )
139 return updates, InfoState(iter_num=state.iter_num + 1)
140
141 return optax.GradientTransformationExtraArgs(init_fn, update_fn)
142
143
144data, data_f = autohf.solve_hf(
145 H,
146 settings=hf_settings | dict(opt_method="grad", steps=-1),
147)
148
149rng = np.random.default_rng(hf_settings["seed"])
150fun = data_f["energy_func"]
151init_params = data["state"]
152init_params += rng.normal(scale=0.001, size=init_params.shape)
153opt = optax.chain(print_info(), optax.lbfgs())
154print(
155 f"Initial value: {fun(init_params):.2e} "
156 f"Initial gradient norm: {otu.tree_l2_norm(jax.grad(fun)(init_params)):.2e}"
157)
158final_params, _ = run_opt(init_params, fun, opt, max_iter=hf_settings["steps"], tol=0.001)
159E_final = fun(final_params)
160print(
161 f"Final value: {E_final:.2e}, "
162 f"Final gradient norm: {otu.tree_l2_norm(jax.grad(fun)(final_params)):.2e}"
163)
164
165
166print("=================")
167print("LBFGS: \t", E_lbfgs)
168print("Grad: \t", E_grad)
169print("custom lbfgs:\t", E_final)