dynamics

Dynamics

📂 dynamics/

• 🐍 __init__.py
• 🐍 dynamics.py
• 📂 pytorch/
  • 🐍 __init__.py
  • 🐍 dynamics.py
• 📂 tensorflow/
  • 🐍 __init__.py
  • 🐍 dynamics.py

Overview

Here we describe the Dynamics object, the main work-horse of our application.

We provide both TensorFlow and PyTorch implementations¹ in pytorch/dynamics.py and tensorflow/dynamics.py

Implementation

We describe below the PyTorch implementation.

from l2hmc.configs import DynamicsConfig
from l2hmc.network.pytorch.network import NetworkFactory

class Dynamics(nn.Module):
    def __init__(
            self,
            potential_fn: Callable,
            config: DynamicsConfig,
            network_factory: Optional[NetworkFactory] = None,
    ) -> None:

Explicitly, our Dynamics object is a subclass of torch.nn.Module that takes as input:

1. potential_fn: Callable:
   The potential ( / action / negative log likelihood) of our theory.

   Should have signature:

   def potential_fn(x: torch.Tensor, beta: float) -> torch.Tensor:

   where beta ~ 1 / Temperature is the inverse coupling constant of our theory.

2. config: DynamicsConfig:
   A @dataclass object defining the configuration used to build a Dynamics object. Looks like:

@dataclass
class DynamicsConfig:
    nchains: int          # num of parallel chains
    group: str            # 'U1' or 'SU3'
    latvolume: List[int]  # lattice volume
    nleapfrog: int        # num leapfrog steps / trajectory
    eps: float            # (initial) step size in leapfrog update
    eps_hmc: Optional[float] = None  # step size for HMC updates
    use_ncp: bool = True  # use Non-Compact Projection for 2D U(1)
    verbose: bool = True
    eps_fixed: bool = False  # use a FIXED step size (non-trainable)
    use_split_xnets: bool = True  # use diff networks for each `x` update
    use_separate_networks: bool = True  # use diff nets for each LF step
    # Switch update style
    # - merge_directions = True:
    #   - N * [forward_update] + N * [backward_update]
    # - merge_directions = False:
    #   - [forward / backward] d ~ U(+,-) ; N * [d update]
    merge_directions: bool = True

Footnotes

1. We strive to keep the implementations as close as possible. ↩