4. Save and Load

4.1. Save Square Hubbard Example

Listing 4.1 square_hubbard_save.py
 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
41output_fname = "autohf_output.pkl"
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": (Ne, Ne),
50  "noncollinear": False,
51  "output": output_fname,
52  "dump_batch": True,
53}
54
55#### Setup calculation
56N = Nx * Ny
57T = makeT(Nx, Ny)
58
59H = autohf.AutoHFHamiltonian(T=(T, T), U=U)
60
61autohf.solve_hf(
62  H,
63  settings=hf_settings,
64)
65
66# the following is an example of what you can do with the output file API
67print("####### Reading output file ######")
68import pickle
69
70with open(output_fname, "rb") as f:
71  output_data = pickle.load(f)
72
73print(f"Best E={output_data['E']}")
74print(f"Ansatz was {output_data['args'].ansatz}")
75orbitals = output_data["orbitals"]
76print("Orbital overlap matrix:")
77Ov = np.swapaxes(orbitals, -1, -2) @ np.conj(orbitals)
78print(np.round(Ov, 4))

4.2. Square Hubbard with Restart Example

Listing 4.2 square_hubbard_with_restart.py
 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 = 8.5
38Ne = 7
39steps = 100
40batch_size = 32
41output_base = "autohf_output"
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": (Ne, Ne),
50  "noncollinear": False,
51  "output": output_base + "_0.pkl",
52  "dump_batch": True,  # needed to convert from SD_ROT to SD
53  "seed": 42,
54}
55
56#### Setup calculation
57N = Nx * Ny
58T = makeT(Nx, Ny)
59
60H = autohf.AutoHFHamiltonian(T=(T, T), U=U)
61
62autohf.solve_hf(
63  H,
64  settings=hf_settings,
65)
66
67
68hf_settings = {
69  "verbose": False,
70  "steps": steps,
71  "opt_method": "lbfgs",
72  "ansatz": "SD",
73  "riemannian": True,
74  "batch_size": batch_size // 2,  # pick the best half
75  "gpu": False,
76  "nelec": (Ne, Ne),
77  "noncollinear": False,
78  "output": output_base + "_1.pkl",
79  "seed": 1234,
80}
81autohf.solve_hf(
82  H,
83  settings=hf_settings,
84  start_from=output_base + "_0.pkl",
85  jaxoptargs=dict(tol=1e-14),  # secret sauce
86)