qcd_ml.nn

qcd_ml.nn.ptc

Parallel Transport Convolutions.

class qcd_ml.nn.ptc.v_PTC(n_feature_in, n_feature_out, paths, U, **path_buffer_kwargs)

Parallel Transport Convolution for objects that transform vector-like.

Weights are stored as [feature_in, feature_out, path].

paths is a list of paths. Every path is a list [(direction, nhops)]. An empty list is the path that does not perform any hops.

For a 1-hop 1-layer model, construct the layer as such:

U = torch.tensor(np.load("path/to/gauge/config.npy"))

paths = [[]] + [[(mu, 1)] for mu in range(4)] + [[(mu, -1)] for mu in range(4)]
layer = v_PTC(1, 1, paths, U)

forward(features_in)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

gauge_transform_using_transformed(U_transformed)

Update the v_PTC layer: The old gauge field U is replaced by U_transformed. The weights are kept.

NOTE: This does not create a transformed copy of the layer!

Instead the layer is updated.

Mostly used for testing.

qcd_ml.nn.lptc

Local Parallel Transport Convolutions.

class qcd_ml.nn.lptc.v_LPTC(n_feature_in, n_feature_out, paths, U, **path_buffer_kwargs)

Local Parallel Transport Convolution for objects that transform vector-like.

Weights are stored as [feature_in, feature_out, path].

paths is a list of paths. Every path is a list [(direction, nhops)]. An empty list is the path that does not perform any hops.

forward(features_in)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

gauge_transform_using_transformed(U_transformed)

Update the v_LPTC layer: The old gauge field U is replaced by U_transformed. The weights are kept.

NOTE: This does not create a transformed copy of the layer!

Instead the layer is updated.

Mostly used for testing.

class qcd_ml.nn.lptc.v_LPTC_NG(n_feature_in, n_feature_out, paths, grid_dims, internal_dof)

Local Parallel Transport Convolution for objects that transform vector-like but with no gauge degrees of freedom.

Weights are stored as [feature_in, feature_out, path].

paths is a list of paths. Every path is a list [(direction, nhops)]. An empty list is the path that does not perform any hops.

forward(features_in)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

pt

Parallel Transport Layers.

class qcd_ml.nn.pt.v_PT(paths, U, **path_buffer_kwargs)

Parallel Transport Layer for objects that transform vector-like.

It has no weights.

paths is a list of paths. Every path is a list [(direction, nhops)]. An empty list is the path that does not perform any hops.

For a PT layer with all 0- and 1-hop paths, construct the layer like this:

U = torch.tensor(np.load("path/to/gauge/config.npy"))

paths = [[]] + [[(mu, 1)] for mu in range(4)] + [[(mu, -1)] for mu in range(4)]
layer = v_PT(paths, U)

forward(features_in)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

gauge_transform_using_transformed(U_transformed)

Update the v_PT layer: The old gauge field U is replaced by U_transformed. The weights are kept.

NOTE: This does not create a transformed copy of the layer!

Instead the layer is updated.

Mostly used for testing.

qcd_ml.nn.dense

This module provides dense linear layers. Currently v_Dense for vector-like objects is provided.

class qcd_ml.nn.dense.v_Dense(n_feature_in, n_feature_out)

Dense Layer for vectors.

v_Dense.forward(features_in) computes

\phi_o(x) = \sum\limits_i W_{io} \phi_i(x)

where W_{io} are spin matrices.

v_Dense.reverse(features_in) computes the hermitian adjoint operation, i.e.,

..math:

\phi_i(x) = \sum\limits_o} W_{io}^\dagger \phi_o(x).

forward(features_in)

\phi_o(x) = \sum\limits_i W_{io} \phi_i(x)

reverse(features_in)

Hermitian adjoint operation of forward.

..math:

\phi_i(x) = \sum\limits_o} W_{io}^\dagger \phi_o(x).

qcd_ml.nn.matrix_layers

Layers for matrix valued fields, i.e., fields that transform as .. math:

M(x) \rightarrow \Omega(x) M(x) \Omega(x).

Provides the following layers:

• LGE_Convolution

• LGE_Bilinear

• LGE_ReTrAct

• LGE_Exp

• PolyakovLoopGenerator and PositiveOrientationPlaquetteGenerator

See [10.1103/PhysRevLett.128.032003].

class qcd_ml.nn.matrix_layers.LGE_CB(n_features_in, n_features_out, paths, disable_cache=True)

A combined lattice gauge equivariant convolution-bilinear layer. Originally described in the supplemental material of https://link.aps.org/doi/10.1103/PhysRevLett.128.032003.

The input features are mapped as

W_{jk}'(x) &= (T_{p_k}W_j)(x) \\ W^a(x) &= (W_j(x), W_j^\dagger(x), \mathbb{1}) \\ W^b(x) &= (W_{jk}'(x), W_{jk}^{\prime\dagger}(x), \mathbb{1})\\ W_i^o(x) &= \sum\limits_{ijj'} \alpha_{ijj'} W_j^a(x) W_{j'}^b(x)

and W^o is returned.

forward(U, features_in)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

---

Activation functions for matrix-like fields, i.e., fields that transform as

M(x) \rightarrow \Omega(x) M(x) \Omega(x)

class qcd_ml.nn.matrix_layers.activation.LGE_ReTrAct(activation, n_features)

Given an activation function activation (F) applies

W_j(x) \rightarrow F(\omega_j \mbox{Re}\mbox{Tr}(W_j(x)) \alpha_j) W_j(x)

forward(features)

W_j(x) \rightarrow F(\omega_j \mbox{Re}\mbox{Tr}(W_j(x)) \alpha_j) W_j(x)

---

This module contains lattice gauge equvariant bilinear layers.

class qcd_ml.nn.matrix_layers.bilinear.Apply_LGE_Bilinear(*args, **kwargs)

static backward(ctx, grad_output)

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, features_in1, features_in2, weights)

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass


@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

class qcd_ml.nn.matrix_layers.bilinear.LGE_Bilinear(n_input1, n_input2, n_output)

This class provides lattice gauge equivariant bilinear layers.

W_{x,i}, W_{x,i}' \rightarrow \sum_{j,k} \alpha_{i,j,k} W_{x,j} W_{x,k}'

See 10.1103/PhysRevLett.128.032003 for more details.

forward(features_in1, features_in2)

This class provides lattice gauge equivariant bilinear layers. .. math:

W_{x,i}, W_{x,i}' \rightarrow \sum_{j,k} \alpha_{i,j,k} W_{x,j} W_{x,k}'

class qcd_ml.nn.matrix_layers.bilinear.LGE_BilinearLM(n_input1, n_input2, n_output)

This is an implementation that provides the same functionality as LGE_Bilinear. The backward pass is slower by a factor of 2, but with significantly less memory consumption.

This class provides lattice gauge equivariant bilinear layers.

W_{x,i}, W_{x,i}' \rightarrow \sum_{j,k} \alpha_{i,j,k} W_{x,j} W_{x,k}'

See 10.1103/PhysRevLett.128.032003 for more details.

forward(features_in1, features_in2)

This class provides lattice gauge equivariant bilinear layers. .. math:

W_{x,i}, W_{x,i}' \rightarrow \sum_{j,k} \alpha_{i,j,k} W_{x,j} W_{x,k}'

---

Convolutions for matrix-like fields, i.e., fields that transform as

M(x) \rightarrow \Omega(x) M(x) \Omega(x)

class qcd_ml.nn.matrix_layers.convolution.LGE_Convolution(n_input, n_output, paths, disable_cache=True)

Provides a convolution for matrix-like fields, i.e., fields that transform as

M(x) \rightarrow \Omega(x) M(x) \Omega(x)

The convolution is defined as

W_i(x) \rightarrow \sum_{j\mu k} \omega_{i\mu k j} U_{\mu k}(x) W_j(x+k\mu) U_{\mu k}^\dagger(x)

See 10.1103/PhysRevLett.128.032003 for more details.

We implement this convolution differently: We define a gauge transporter along an arbitrary path T_p as

(T_p(M))(x) = ((\prod\limits_{\mu_k \in p} H_{\mu_k}) M)(x)

where

(H_{\mu} M)(x) = U_{\mu}(x) M(x + \mu) U_{\mu}(x+\mu)^\dagger

Then, the convolution is defined as

W_i(x) \rightarrow \sum_{jik} \omega_{j i k} T_{p_k}(W_j)(x)

clear_path_buffers()

If disable_cache=False, this method can be used to clear the pre-computed cache.

forward(U, features_in)

W_i(x) \rightarrow \sum_{j\mu k} \omega_{i\mu k j} U_{\mu k}(x) W_j(x+k\mu) U_{\mu k}^\dagger(x)

---

class qcd_ml.nn.matrix_layers.exponentiation.LGE_Exp(n_features_in, matrix_mode='a')

Provides an exponentiation layer for matrix-like fields acting on gauge links, i.e.,

U_\mu(x) \rightarrow \exp\left(\sum\limits_j \beta_{\mu,i} W_{i}(x)\right) U_\mu(x)

forward(U, W)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

---

This module provides loop generator layers, i.e., layers that take a link field U and compute a set of gauge equivariant loops.

All layers inherit from AbstractLoopGenerator and provide the class property nfeatures_out. This allows to access the number of generated loops programatically:

features_in = PositiveOrientationPlaquetteGenerator.nfeatures_out
generator = PositiveOrientationPlaquetteGenerator()
...

class qcd_ml.nn.matrix_layers.loop_generator.PolyakovLoopGenerator(disable_cache=True)

Generates the Polyakov loops

P_\mu(x) = \prod\limits_{k=0}^{L_\mu} U_\mu(x + k\mu)

for \mu = 0,1,2,3.

forward(U)

Compite the Polyakov loops.

class qcd_ml.nn.matrix_layers.loop_generator.PositiveOrientationPlaquetteGenerator(disable_cache=True)

Generates the positivly oriented plaquettes, i.e.,

P_{\mu\nu}(x) = U_\mu(x) U_\nu(x+\mu) U_\mu^\dagger(x+\nu) U_\nu^\dagger(x)

forward(U)

Compute the positivly oriented plaquettes.