initial release version

FD.py

@@ -0,0 +1,221 @@
+"""Finite Difference"""
+import numpy as np
+from numpy import *
+from numpy.linalg import *
+from functools import reduce
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.nn.functional as F
+from aTEAM.nn.modules import MK
+from aTEAM.nn.functional import periodicpad
+
+__all__ = ['MomentBank','FD1d','FD2d','FD3d']
+
+def _inv_equal_order_m(d,m):
+    A = []
+    assert d >= 1 and m >= 0
+    if d == 1:
+        A = [[m,],]
+        return A
+    if m == 0:
+        for i in range(d):
+            A.append(0)
+        return [A,]
+    for k in range(m+1):
+        B = _inv_equal_order_m(d-1,m-k)
+        for b in B:
+            b.append(k)
+        A = A+B
+    return A
+
+def _less_order_m(d,m):
+    A = []
+    for k in range(m+1):
+        B = _inv_equal_order_m(d,k)
+        for b in B:
+            b.reverse()
+        B.sort()
+        B.reverse()
+        A.append(B)
+    return A
+
+def _torch_setter_by_index(t,i,v):
+    for j in i[:-1]:
+        t = t[j]
+    t[i[-1]] = v
+def _torch_reader_by_index(t,i):
+    for j in i:
+        t = t[j]
+    return t
+
+class MomentBank(nn.Module):
+    """
+    generate moment matrix bank for differential kernels with order 
+    no more than max_order.
+    Arguments:
+        dim (int): dimension
+        kernel_size (tuple of int): size of differential kernels
+        max_order (int): max order of differential kernels
+        dx (double): the MomentBank.kernel will automatically compute kernels 
+            according to MomentBank.moment and MomentBank.dx
+        constraint (string): 'moment' or 'free'. Determine MomentBank.x_proj 
+            and MomentBank.grad_proj
+    """
+    def __init__(self, dim, kernel_size, max_order, dx=1.0, constraint='moment'):
+        super(MomentBank, self).__init__()
+        self._dim = dim
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size,]*self.dim
+        assert min(kernel_size) > max_order
+        self.m2k = MK.M2K(kernel_size)
+        self._kernel_size = kernel_size.copy()
+        self._max_order = max_order
+        if not iterable(dx):
+            dx = [dx,]*dim
+        self._dx = dx.copy()
+        self.constraint = constraint
+        d = dim
+        m = max_order
+        self._order_bank = _less_order_m(d, m)
+        N = 0
+        for a in self._order_bank:
+            N += len(a)
+        moment = torch.DoubleTensor(*([N,]+kernel_size)).zero_()
+        index = zeros([m+1,]*dim,dtype=np.int64)
+        for i,o in enumerate(self.flat_order_bank()):
+            _torch_setter_by_index(moment[i],o,1)
+            _torch_setter_by_index(index,o,i)
+            # moment[i,*o] = 1
+            # index[*o] = i
+        self.moment = nn.Parameter(moment)
+        self._index = index
+        scale = torch.from_numpy(ones((self.moment.size()[0])))
+        l = lambda a,b:a*b
+        for i,o in enumerate(self.flat_order_bank()):
+            s = reduce(l, (self.dx[j]**oj for j,oj in enumerate(o)), 1)
+            scale[i] = 1/s
+        self.register_buffer('scale',scale)
+
+    def __index__(self,*args):
+        return self.moment[_torch_reader_by_index(self._index, args)]
+
+    def dim(self):
+        return self._dim
+    @property
+    def dx(self):
+        return self._dx.copy()
+
+    def kernel(self):
+        scale = Variable(self.scale[:,newaxis])
+        kernel = self.m2k(self.moment)
+        size = kernel.size()
+        kernel = kernel.view([size[0],-1])
+        return (kernel*scale).view(size)[:,newaxis]
+
+    def flat_order_bank(self):
+        for a in self._order_bank:
+            for o in a:
+                yield o
+    def _proj_(self,M,s,c):
+        for j in range(s):
+            for o in self._order_bank[j]:
+                _torch_setter_by_index(M,o,c)
+                # M[*o] = c
+    def x_proj(self,*args,**kw):
+        if self.constraint == 'free':
+            return None
+        if isinstance(self.constraint,int):
+            acc = self.constraint
+        else:
+            acc = 1
+        for i,o in enumerate(self.flat_order_bank()):
+            self._proj_(self.moment.data[i],sum(o)+acc,0)
+            _torch_setter_by_index(self.moment.data[i], o, 1)
+            # self.moment.data[i,*o] = 1
+        return None
+    def grad_proj(self,*args,**kw):
+        if self.constraint == 'free':
+            return None
+        if isinstance(self.constraint,int):
+            acc = self.constraint
+        else:
+            acc = 1
+        for i,o in enumerate(self.flat_order_bank()):
+            self._proj_(self.moment.grad.data[i],sum(0)+acc,0)
+        return None
+
+    def forward(self):
+        return self.kernel()
+#%%
+
+class _FDNd(nn.Module):
+    """
+    Finite difference automatically handle boundary conditions
+    Arguments for class:`_FDNd`:
+        dim (int): dimension
+        kernel_size (tuple of int): finite difference kernel size
+        boundary (string): 'Dirichlet' or 'Periodic'
+    Arguments for class:`MomentBank`:
+        max_order, dx, constraint
+    """
+    def __init__(self, dim, kernel_size, boundary='Dirichlet'):
+        super(_FDNd, self).__init__()
+        self._dim = dim
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size,]*self.dim
+        self._kernel_size = kernel_size.copy()
+        padwidth = []
+        for k in reversed(kernel_size):
+            padwidth.append((k-1)//2)
+            padwidth.append(k-1-(k-1)//2)
+        self._padwidth = padwidth
+        self.boundary = boundary.upper()
+
+    def dim(self):
+        return self._dim
+    @property
+    def padwidth(self):
+        return self._padwidth.copy()
+    @property
+    def boundary(self):
+        return self._boundary
+    @boundary.setter
+    def boundary(self,v):
+        self._boundary = v.upper()
+    def pad(self, inputs):
+        if self.boundary == 'DIRICHLET':
+            return F.pad(inputs, self.padwidth)
+        else:
+            return periodicpad(inputs, self.padwidth)
+
+    def conv(self, inputs, weight):
+        raise NotImplementedError
+    def forward(self, inputs, kernel):
+        """
+        Arguments:
+            inputs (Variable): torch.size: (batch_size, spatial_size[0], spatial_size[1], ...)
+        """
+        inputs = self.pad(inputs)
+        inputs = inputs[:,newaxis]
+        return self.conv(inputs, kernel)
+
+class FD1d(_FDNd):
+    def __init__(self, kernel_size, max_order, dx=1.0, constraint='moment', boundary='Dirichlet'):
+        super(FD1d, self).__init__(1, kernel_size, boundary=boundary)
+        self.MomentBank = MomentBank(1, kernel_size, max_order, dx=dx, constraint=constraint)
+        self.conv = F.conv1d
+        # self.kernel = self.MomentBank.kernel
+class FD2d(_FDNd):
+    def __init__(self, kernel_size, max_order, dx=1.0, constraint='moment', boundary='Dirichlet'):
+        super(FD2d, self).__init__(2, kernel_size, boundary=boundary)
+        self.MomentBank = MomentBank(2, kernel_size, max_order, dx=dx, constraint=constraint)
+        self.conv = F.conv2d
+        # self.kernel = self.MomentBank.kernel
+class FD3d(_FDNd):
+    def __init__(self, kernel_size, max_order, dx=1.0, constraint='moment', boundary='Dirichlet'):
+        super(FD3d, self).__init__(3, kernel_size, boundary=boundary)
+        self.MomentBank = MomentBank(3, kernel_size, max_order, dx=dx, constraint=constraint)
+        self.conv = F.conv3d
+        # self.kernel = self.MomentBank.kernel
+

checkpoint/linpde.yaml

@@ -0,0 +1,24 @@
+--batch_size: 24
+--constraint: moment # constraint type: can be moment\frozen\free
+--dt: 0.015 # time step size for each delta-t block
+--end_noise_level: 0.015 # noise level for generated data at t=T
+--gpu: 0
+--initfreq: 4
+--interp_degree: 2
+--interp_mesh_size: 20
+--kernel_size: 5 # convolution kernel size: 5 or 7
+--layer: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
+--max_order: 4
+--maxiter: 20000
+--precision: double # precision of modules: can be double\float
+--recordcycle: 200
+--recordfile: convergence # convergence record file
+--repeatnum: 25
+--savecycle: 10000
+--start_noise_level: 0.015 # noise level for generated data at t=0
+--taskdescriptor: linpde5x5moment4order0.015dt0.015noise-double
+--teststepnum: 80
+--variant_coe_magnitude: 1.0
+--xn: 50
+--yn: 50
+

checkpoint/nonlinpde.yaml

@@ -0,0 +1,28 @@
+--batch_size: 24
+--constraint: moment # constraint type: can be moment\frozen\free
+--diffusivity: 0.3
+--dt: 0.01 # time step size for each delta-t block
+--end_noise_level: 0.01 # noise level for generated data at t=T
+--gpu: 0
+--initfreq: 4
+--interp_degree: 2
+--interp_mesh_size: 5
+--kernel_size: 7 # convolution kernel size: 5 or 7
+--layer: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
+--max_order: 2
+--maxiter: 20000
+--nonlinear_coefficient: 15
+--nonlinear_interp_degree: 4
+--nonlinear_interp_mesh_bound: 15.0
+--nonlinear_interp_mesh_size: 20
+--precision: double # precision of modules: can be double\float
+--recordcycle: 200
+--recordfile: convergence # convergence record file
+--repeatnum: 25
+--savecycle: 10000
+--start_noise_level: 0.01 # noise level for generated data at t=0
+--taskdescriptor: nonlinpde7x7moment2order-double
+--teststepnum: 80
+--xn: 50
+--yn: 50
+

errs_compare.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+compare relative errs
+"""
+#%% 
+from numpy import *
+import torch
+import matplotlib.pyplot as plt
+errs = []
+errs.append(torch.load('checkpoint/linpde5x5moment4order0.015dt0.015noise-double/errs.pkl')) # blue
+errs.append(torch.load('checkpoint/linpde7x7moment4order0.015dt0.015noise-double/errs.pkl')) # orange
+
+edgecolorlist = ('#1B2ACC','#CC4F1B') # , 'red', 'yellow')
+facecolorlist = ('#089FFF','#FF9848') # , 'red', 'yellow')
+
+alpha = 0.7 # facecolor transparency
+
+showlayer = [0,7,10,15]
+fig,ax = plt.subplots(1,len(showlayer), sharex=False, sharey=True)
+title = ''
+upq = 75
+downq = 25
+n = 80
+x = arange(1,n,dtype=float64)
+j = 0
+i = 0
+for l in showlayer:
+    ll = l
+    if l == 0:
+        ll = 'warmup'
+    else:
+        ll = 'layer-'+str(l)
+    for s in range(len(edgecolorlist)):
+        y = errs[s][l][:,1:n]
+        y_mean = sqrt(y).mean(axis=1)
+        y_up = percentile(sqrt(y),q=upq,axis=0)
+        y_down = percentile(sqrt(y),q=downq,axis=0)
+        ax.flatten()[j].fill_between(x,y_down,y_up,edgecolor=edgecolorlist[s], facecolor=facecolorlist[s],\
+                linestyle='-', alpha=alpha)
+    if l == 0:
+        ax.flatten()[j].set_title(r'warm-up', fontsize=20)
+    else:
+        ax.flatten()[j].set_title(r''+str(l)+' $\delta t$-block', fontsize=20)
+    ax.flatten()[j].set_yscale('log')
+    ax.flatten()[j].set_ylim(1e-2,1e2)
+    ax.flatten()[j].set_xticks([1,20,40,60])
+    ax.flatten()[j].grid()
+    j += 1
+
+#%%
+
+

learn_linpde.sh

@@ -0,0 +1,10 @@
+PRECISION="double"
+CONSTRAINT="moment"
+KERNEL_SIZE=5
+MAX_ORDER=4
+DT=0.015
+NOISE_LEVEL=0.015
+TASKDESCRIPTOR="linpde"${KERNEL_SIZE}x${KERNEL_SIZE}${CONSTRAINT}${MAX_ORDER}order${DT}dt${NOISE_LEVEL}noise-${PRECISION}
+echo ${TASKDESCRIPTOR}
+python learn_variantcoelinear2d.py --precision=${PRECISION} --taskdescriptor=${TASKDESCRIPTOR} --constraint=${CONSTRAINT} --kernel_size=${KERNEL_SIZE} --max_order=${MAX_ORDER} --dt=${DT} --start_noise_level=${NOISE_LEVEL} --end_noise_level=${NOISE_LEVEL}
+python linpdetest.py ${TASKDESCRIPTOR}

learn_nonlinpde.sh

@@ -0,0 +1,8 @@
+PRECISION="double"
+CONSTRAINT="moment"
+KERNEL_SIZE=7
+MAX_ORDER=2
+TASKDESCRIPTOR="nonlinpde"${KERNEL_SIZE}x${KERNEL_SIZE}${CONSTRAINT}${MAX_ORDER}order-${PRECISION}
+echo ${TASKDESCRIPTOR}
+python learn_singlenonlinear2d.py --kernel_size=${KERNEL_SIZE} --max_order=${MAX_ORDER} --precision=${PRECISION} --constraint=${CONSTRAINT} --taskdescriptor=${TASKDESCRIPTOR}
+python nonlinpdetest.py ${TASKDESCRIPTOR}