import mylatlib

L, T = 64, 128
gauge_field = mylatlib.cold_start(L, T)

gauge_field.heatbath(gauge_action="wilson", num_updates=100)

correlators = []

for i in range(100):
    
    prop = mylatlib.compute_wilson_prop(mass=0.1, gauge_field=gauge_field)
    correlators.append(mylatlib.meson256(prop, prop))
    
    gauge_field.heatbath(gauge_action="wilson", num_updates=10)
    
mylatlib.save("the_results", correlators)

import pyQCD

pyQCD.Lattice?

lattice = pyQCD.Lattice(L=4, T=8,
                        beta=5.5, action="wilson",
                        meas_spacing=10,
                        update_method="heatbath")

for i in range(20):
    print(lattice.get_av_plaquette())
    lattice.update()

config = lattice.get_config()

print("Config type: {}".format(type(config)))
print("Config shape: {}".format(config.shape))
print("")
print("Gauge field for link (t, x, y, z, mu) = (4, 3, 1, 2, 2):")
print(config[4, 3, 1, 2, 2])

lattice = pyQCD.Lattice(8, 4, 5.8, "rectangle_improved", 10)
# Thermalize with 100 updates
lattice.thermalize(100)

for i in range(10):
    print("Generating config {}".format(i))
    # Do ten updates
    lattice.next_config()
    lattice.save_config("config{}.npy".format(i))

from functools import partial

lattice = pyQCD.Lattice(4, 8, 5.5, "wilson", 10)
inverter = partial(lattice.invert_wilson_dirac, mass=0.2)
source = lattice.point_source([0, 0, 0, 0])

print("Source type: {}".format(type(source)))
print("Source shape: {}".format(source.shape))
print("Source non-zero entries: {}".format(source.nonzero()))

%matplotlib inline

import logging
import matplotlib.pyplot as plt

logger = logging.getLogger()
logger.setLevel(logging.INFO)

prop = pyQCD.compute_propagator(source, inverter)
print("")
print("Propagator shape: {}".format(prop.shape))
ps_correlator = pyQCD.compute_meson_corr(prop, prop, "g5", "g5")
plt.plot(ps_correlator)
plt.xlabel("$t$", fontsize=18)
plt.ylabel("$C(t)$", fontsize=18)
plt.show()

# Specify the lattice and simulation
lattice = pyQCD.Lattice(16, 32, 4.41, "rectangle_improved", 10)
sim = pyQCD.Simulation(lattice, num_configs=100, num_warmup_updates=100)
# Specify the callback function to load a config from disk during the simulation
def load_config(index):
    # N.B. These configs don't exist, so this simulation won't run.
    return pyQCD.io.load_ildg_config("some_config.lime{}".format(index))
# Tell the simulation 
sim.specify_ensemble(load_config, range(10, 1010, 10))

@pyQCD.Log(ignore=("lattice,")) # Decorator logs the function arguments
def compute_correlators(lattice, mass1, mass2):
    """Compute all 256 meson correlators for the supplied masses"""
    invert_wilson = lattice.invert_wilson_dirac
    # Specify inversion functions
    inverter1 = partial(invert_wilson, mass=mass1, precondition=True)
    inverter2 = partial(invert_wilson, mass=mass2, precondition=True)
    # Define a point source
    src = lattice.point_source([0, 0, 0, 0])
    # Compute our propagators
    prop1 = pyQCD.compute_propagator(src, inverter1)
    prop2 = (prop1 if mass1 == mass2
             else pyQCD.compute_propagator(src, inverter2))
    # Use the props to compute meson correlators
    return pyQCD.compute_meson_corr256(prop1, prop2)

# Add the measurement, specifying the function to write the result to disk
sim.add_measurement(compute_correlators,
                    pyQCD.io.write_datum_callback("16c32_m0.2_m0.2.zip"),
                    kwargs={'mass1': 0.2, 'mass2': 0.2})

sim.run()

import numpy as np
import scipy.sparse.linalg as spla

logger.setLevel(logging.WARN)

lattice = pyQCD.Lattice(4, 8, 5.5, "wilson", 10)

def matvec(psi):
    return lattice.apply_wilson_dirac(psi, 0.2).flatten()

N = 12 * np.prod(lattice.shape)
linop = spla.LinearOperator(shape=(N, N), matvec=matvec)

eigvals, eigvecs = spla.eigs(linop, k=200)

plt.plot(eigvals.real, eigvals.imag, 'o')
plt.xlabel("$\mathrm{Re}\lambda$", fontsize=18)
plt.ylabel("$\mathrm{Im}\lambda$", fontsize=18)
plt.show()

# Example input file: wilson.py

import numpy as np

import pyQCD

@types(mass_="double", hoppingMatrix_='HoppingTerm*')
class Wilson(object):

    @types(mass="const double")
    def __init__(self, mass):

        self.mass_ = mass
        self.hoppingMatrix_ = HoppingTerm(1)

    @types(psi='VectorXcd', eta='VectorXcd')
    def apply(self, psi):

        eta = (4 + self.mass_) * psi
        eta -= 0.5 * self.hoppingMatrix_.apply(psi)
        return eta

pyQCD.codegen.gen_from_src("wilson.py")

cat wilson.cpp