add: hmc, phi4, gaussian, roll

Author: li012589

.gitignore

@@ -0,0 +1,8 @@
+.cache
+**/*.*~
+**/*.pyc
+**/.DS_Store
+**/test.py
+**/*.testSave
+docs/_build
+data

flow/__init__.py

@@ -0,0 +1 @@
+from .flow import Flow

flow/flow.py

@@ -0,0 +1,36 @@
+import numpy as np
+import torch
+from torch import nn
+
+class Flow(nn.Module):
+
+    def __init__(self, prior = None,name = "Flow"):
+        super(Flow, self).__init__()
+        self.name = name
+        self.prior = prior
+
+    def __call__(self,*args,**kargs):
+        return self.sample(*args,**kargs)
+
+    def sample(self,batchSize):
+        raise NotImplementedError(str(type(self)))
+
+    def inference(self,x):
+        raise NotImplementedError(str(type(self)))
+
+    def generate(self,z):
+        raise NotImplementedError(str(type(self)))
+
+    def logProbability(self,x):
+        z,logp = self.inference(x)
+        if self.prior is not None:
+            return self.prior.logProbability(z)+logp
+        return logp
+
+    def save(self,saveDict):
+        saveDict[self.name] = self.state_dict()
+        return saveDict
+
+    def load(self,saveDict):
+        self.load_state_dict(saveDict)
+        return saveDict

model/__init__.py flow/hierarchy/__init__.py

@@ -0,0 +1,11 @@
+import torch
+from torch import nn
+
+from model import Flow
+
+class RNVP(Flow):
+    def __init__(self, maskList, tList, sList, prior = None, name = "RNVP"):
+        super(RNVP,self).__init__(prior,name)
+        self.maskList = nn.Parameter(maskList)
+        self.tList = tList
+        self.sList = sList

model/hierarchy/mera.py flow/hierarchy/mera.py

@@ -0,0 +1,2 @@
+from .gaussian import Gaussian
+from .phi4 import Phi4

model/hierarchy/tebd.py flow/hierarchy/tebd.py

@@ -0,0 +1,16 @@
+import torch
+import numpy as np
+
+from .source import Source
+
+class Gaussian(Source):
+    def __init__(self, nvars, sigma = 1, name="gaussian", requiresGrad = False):
+        super(Gaussian,self).__init__(nvars,name)
+        self.sigma = torch.nn.Parameter(torch.tensor([sigma],dtype=torch.float32),requires_grad=requiresGrad)
+
+    def sample(self, batchSize):
+        size = [batchSize] + self.nvars
+        return torch.randn(size).to(self.sigma)
+
+    def energy(self, z):
+        return -(-0.5 * (z/self.sigma)**2-0.5*torch.log(2.*np.pi*self.sigma**2)).view(z.shape[0],-1).sum(dim=1)

model/hierarchy/template.py flow/hierarchy/template.py

@@ -0,0 +1,41 @@
+import numpy as np
+import torch
+
+from .source import Source
+from utils import HMC
+from utils import roll
+
+class Phi4(Source):
+    def __init__(self,l,dims,kappa,lamb,name = None):
+        if name is None:
+            self.name = "phi4_l"+str(l)+"_d"+str(dims)+"_kappa"+str(kappa)+"_lamb"+str(lamb)
+        else:
+            self.name = name
+        self.kappa = kappa
+        self.lamb = lamb
+        self.dims = dims
+        nvars = []
+        for _ in range(dims):
+            nvars += [l]
+        super(Phi4,self).__init__(nvars,name)
+
+    def sample(self, batchSize, thermalSteps = 50, interSteps=5, epsilon=0.1):
+        inital = torch.randn([batchSize]+self.nvars,requires_grad=True)
+        inital = HMC(self.energy,inital,thermalSteps,interSteps,epsilon)
+        return inital.detach()
+
+    def energy(self,x):
+        S = 0
+        for i in range(self.dims):
+            S += x*roll(x,[1],[i+1])
+            #S += x*roll(x,[-1],[i+1])
+        term1 = x**2
+        term2 = (term1-1)**2
+        for _ in range(self.dims):
+            S = S.sum(-1)
+            term1 = term1.sum(-1)
+            term2 = term2.sum(-1)
+        S *= -2*self.kappa
+        term2 *= self.lamb
+        S += term1 + term2
+        return -S

model/hierarchy/__init__.py flow/rnvp/__init__.py

@@ -0,0 +1,30 @@
+import numpy as np
+import torch
+from torch import nn
+
+class Source(nn.Module):
+
+    def __init__(self, nvars,name = "Flow"):
+        super(Source, self).__init__()
+        self.name = name
+        self.nvars = nvars
+
+    def __call__(self,*args,**kargs):
+        return self.sample(*args,**kargs)
+
+    def sample(self,batchSize):
+        raise NotImplementedError(str(type(self)))
+
+    def logProbability(self,x):
+        return -self.energy(x)
+
+    def energy(self,x):
+        raise NotImplementedError(str(type(self)))
+
+    def save(self,saveDict):
+        saveDict[self.name] = self.state_dict()
+        return saveDict
+
+    def load(self,saveDict):
+        self.load_state_dict(saveDict)
+        return saveDict

flow/rnvp/rnvp.py

@@ -0,0 +1,2 @@
+from .mc import HMCsampler, HMC
+from .roll import roll

model/rnvp/rnvp.py

@@ -0,0 +1 @@
+from .hmc import HMCsampler, HMC

model/template.py

@@ -0,0 +1,47 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.autograd import grad as torchgrad
+
+torch.manual_seed(42)
+
+def HMC(energy,x,length,steps,epsilon):
+    def grad(z):
+        return torchgrad(energy(z),z,grad_outputs=torch.ones(z.shape[0]))[0]
+
+    E = energy(x)
+    g = grad(x)
+
+    for l in range(length):
+        p = x.new_empty(size=x.size()).normal_()
+        H = (0.5*p*p).view(p.shape[0], -1).sum(dim=1) + E
+        xnew = x
+        gnew = g
+        for _ in range(steps):
+            p = p- epsilon* gnew/2.
+            xnew = xnew + epsilon * p
+            gnew = grad(xnew)
+            p = p- epsilon* gnew/2.
+        Enew = energy(xnew)
+        Hnew = (0.5*p*p).view(p.shape[0], -1).sum(dim=1) + Enew
+        diff = H-Hnew
+        accept = (diff.exp() >= diff.uniform_()).to(x)
+
+        E = accept*Enew + (1.-accept)*E
+        accept = accept.view(x.shape[0], 1, 1)
+        x = accept*xnew + (1.-accept)*x
+        g = accept*gnew + (1.-accept)*g
+
+    return x
+
+
+class HMCsampler(nn.Module):
+    def __init__(self,energy,nvars, epsilon=0.01, interSteps=10 , thermalSteps = 10):
+        super(HMCsampler,self).__init__()
+        self.nvars = nvars
+        self.energy = energy
+        self.interSteps = interSteps
+        self.inital = HMC(self.energy,torch.randn(nvars),thermalSteps,interSteps)
+
+    def step(self):
+        return HMC(self.energy,self.inital,1,interSteps,epsilon)

source/__init__.py

@@ -0,0 +1,13 @@
+import torch
+
+def roll(x, step, axis):
+    shape = x.shape
+    for i,s in enumerate(step):
+        if s >=0:
+            x1 = x.narrow(axis[i],0,s)
+            x2 = x.narrow(axis[i],s,shape[axis[i]]-s)
+        else:
+            x2 = x.narrow(axis[i],shape[axis[i]]+s,-s)
+            x1 = x.narrow(axis[i],0,shape[axis[i]]+s)
+        x = torch.cat([x2,x1],axis[i])
+    return x