1. Basics

1.1. Pinned Cylinder Example

Listing 1.1 pinned_cyl.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 makeTPin(Nx, Ny, h_pin, yperiodic=False):
17  N = Nx * Ny
18  T = np.zeros((N, N))
19  for i in range(Nx):
20    for j in range(Ny):
21      a = i * Ny + j
22      if i + 1 < Nx or yperiodic:
23        b = ((i + 1) % Nx) * Ny + j
24        T[a, b] = -1
25        T[b, a] = -1
26      b = i * Ny + (j + 1) % Ny
27      T[a, b] = -1
28      T[b, a] = -1
29  Sz = 1 / 2 * h_pin
30  pin = np.zeros(N)
31  pin[np.arange(N - Ny, N)] = Sz * (-1) ** (np.arange(Ny))
32  pin[np.arange(Ny)] = Sz * (-1) ** (1 + np.arange(Ny))
33  print(pin)
34  print(np.sum(T, axis=0))
35  assert np.allclose(T, T.T)
36  return T, pin
37
38
39h_pin = 0.5  # h_pin*Sz_i = h_pin/2*Z
40Nx, Ny = 4, 4
41N = Nx * Ny
42T_, pin = makeTPin(Nx, Ny, h_pin)
43T_up, T_dn = T_ + np.diag(pin), T_ - np.diag(pin)
44
45H = autohf.AutoHFHamiltonian(T=(T_up, T_dn), U=1.5)
46
47hf_settings = {
48  "verbose": False,
49  "steps": 100,
50  "opt_method": "lbfgs",
51  "ansatz": "SD_ROT",
52  "batch_size": 4,
53  "gpu": False,
54  "nelec": (7, 7),
55}
56autohf.solve_hf(
57  H,
58  settings=hf_settings,
59)

1.2. Square Hubbard Example

Listing 1.2 square_hubbard.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
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": (Ne, Ne),
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
59autohf.solve_hf(
60  H,
61  settings=hf_settings,
62)

1.3. Triangular Hubbard GHF Example

Listing 1.3 triangular_hubbard_ghf.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
16hops = (np.array([ 0,  0,  0,  0,  0,  0,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,
17                        2,  3,  3,  3,  3,  3,  3,  4,  4,  4,  4,  4,  4,  5,  5,  5,  5,
18                        5,  5,  6,  6,  6,  6,  6,  6,  7,  7,  7,  7,  7,  7,  8,  8,  8,
19                        8,  8,  8,  9,  9,  9,  9,  9,  9, 10, 10, 10, 10, 10, 10, 11, 11,
20                       11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 14,
21                       14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16, 16, 16,
22                       17, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19,
23                       19, 20, 20, 20, 20, 20, 20, 21, 21, 21, 21, 21, 21, 22, 22, 22, 22,
24                       22, 22, 23, 23, 23, 23, 23, 23, 24, 24, 24, 24, 24, 24, 25, 25, 25,
25                       25, 25, 25, 26, 26, 26, 26, 26, 26, 27, 27, 27, 27, 27, 27, 28, 28,
26                       28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 30, 30, 30, 30, 30, 30, 31,
27                       31, 31, 31, 31, 31, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33, 33,
28                       34, 34, 34, 34, 34, 34, 35, 35, 35, 35, 35, 35]),
29         np.array([ 1,  5,  6, 11, 30, 31,  0,  2,  6,  7, 31, 32,  1,  3,  7,  8, 32,
30                        33,  2,  4,  8,  9, 33, 34,  3,  5,  9, 10, 34, 35,  0,  4, 10, 11,
31                        30, 35,  0,  1,  7, 11, 12, 17,  1,  2,  6,  8, 12, 13,  2,  3,  7,
32                         9, 13, 14,  3,  4,  8, 10, 14, 15,  4,  5,  9, 11, 15, 16,  0,  5,
33                         6, 10, 16, 17,  6,  7, 13, 17, 18, 23,  7,  8, 12, 14, 18, 19,  8,
34                         9, 13, 15, 19, 20,  9, 10, 14, 16, 20, 21, 10, 11, 15, 17, 21, 22,
35                         6, 11, 12, 16, 22, 23, 12, 13, 19, 23, 24, 29, 13, 14, 18, 20, 24,
36                        25, 14, 15, 19, 21, 25, 26, 15, 16, 20, 22, 26, 27, 16, 17, 21, 23,
37                        27, 28, 12, 17, 18, 22, 28, 29, 18, 19, 25, 29, 30, 35, 19, 20, 24,
38                        26, 30, 31, 20, 21, 25, 27, 31, 32, 21, 22, 26, 28, 32, 33, 22, 23,
39                        27, 29, 33, 34, 18, 23, 24, 28, 34, 35,  0,  5, 24, 25, 31, 35,  0,
40                         1, 25, 26, 30, 32,  1,  2, 26, 27, 31, 33,  2,  3, 27, 28, 32, 34,
41                         3,  4, 28, 29, 33, 35,  4,  5, 24, 29, 30, 34])
42        )  # fmt: skip
43
44
45#### Define Parameters
46Nx, Ny = 6, 6
47U = 8
48Ne = 18
49steps = 2000
50batch_size = 4
51
52hf_settings = {
53  "verbose": False,
54  "steps": steps,
55  "opt_method": "lbfgs",
56  "ansatz": "SD_ROT",
57  "batch_size": batch_size,
58  "gpu": False,
59  "nelec": (Ne, Ne),
60  "noncollinear": True,
61  "state0_scale": 0.001,  # start close to U=0
62}
63
64#### Setup calculation
65N = Nx * Ny
66assert Nx == 6 and Ny == 6  # harded coded here
67T = np.zeros((N, N))
68T[hops] = -1
69
70H = autohf.AutoHFHamiltonian(T=(T, T), U=U)
71
72autohf.solve_hf(
73  H,
74  settings=hf_settings,
75)