def test_bijective():
p = source.Gaussian([4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.view(1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
repeat = 2
t = flow.MERA(2, length, layers, repeat, p)
def op(x):
return -x
sym = [op]
m = train.Symmetrized(t, sym)
z = m.prior.sample(100)
xz1, _ = m.inverse(z)
xz2, _ = m.inverse(z)
p1 = m.logProbability(xz1)
p2 = m.logProbability(xz2)
z1, _ = m.forward(xz1)
xz1p, _ = m.inverse(z1)
assert ((xz1 == xz2).sum() + (xz1 == -xz2).sum()) == 100 * 4 * 4
assert_array_almost_equal(p1.detach().numpy(),
p2.detach().numpy(),
decimal=5)
assert_array_almost_equal(np.fabs(xz1.detach().numpy()),
np.fabs(xz1p.detach().numpy()))
def extractFlow(flowCon):
from copy import deepcopy
layers = []
_rnvp = deepcopy(flowCon.layerList[1].flow)
_diag = flowCon.layerList[0]
nn = _diag.shift.shape[0] // 2
_op = source.Gaussian([nn]).to(torch.float64)
assert _diag.shift.sum() == 0
assert _diag.fix.sum() == nn
layers.append(
flow.DiagScaling(nn, initValue=_diag.elements.clone().detach()[:nn]))
layers.append(_rnvp)
return flow.FlowNet(layers, _op).double()
def symmetryMERAInit(L,
d,
nlayers,
nmlp,
nhidden,
nrepeat,
symmetryList,
device,
dtype,
name=None):
s = source.Gaussian([L] * d)
depth = int(math.log(L, 2)) * nrepeat * 2
MaskList = []
for _ in range(depth):
masklist = []
for n in range(nlayers):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.view(1, 2, 2)
else:
b = 1 - b
masklist.append(b)
masklist = torch.cat(masklist, 0).to(torch.float32)
MaskList.append(masklist)
dimList = [4]
for _ in range(nmlp):
dimList.append(nhidden)
dimList.append(4)
layers = [
flow.RNVP(MaskList[n], [
utils.SimpleMLPreshape(dimList, [nn.ELU()
for _ in range(nmlp)] + [None])
for _ in range(nlayers)
], [
utils.SimpleMLPreshape(
dimList, [nn.ELU()
for _ in range(nmlp)] + [utils.ScalableTanh(4)])
for _ in range(nlayers)
]) for n in range(depth)
]
f = flow.MERA(2, L, layers, nrepeat, s)
f = Symmetrized(f, symmetryList, name=name)
f.to(device=device, dtype=dtype)
return f
def test_bijective():
p = source.Gaussian([2,2])
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(1,4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:,i] = 1
b=b.reshape(1,2,2)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
f = flow.RNVP(maskList, [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),utils.ScalableTanh(4)]) for _ in range(4)],p)
bijective(f)
def test_symplectic():
p = source.Gaussian([8])
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[i] = 1
b=b.reshape(1,4)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
fl = flow.RNVP(maskList, [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),utils.ScalableTanh(4)]) for _ in range(4)])
f = flow.PointTransformation(fl,p)
symplectic(f)
def test_bijective():
p = source.Gaussian([4])
f1 = flow.Scaling(4, [2, 3])
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 4)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
f2 = flow.NICE(maskList, [
utils.SimpleMLPreshape([4, 32, 32, 4],
[nn.ELU(), nn.ELU(), None]) for _ in range(4)
])
f = flow.FlowNet([f1, f2], p)
bijective(f)
def test_bijective():
p = source.Gaussian([4,4])
BigList = []
for _ in range(2*2*2):
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(1,4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:,i] = 1
b=b.view(1,2,2)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
BigList.append(maskList)
layers = [flow.RNVP(BigList[n], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),utils.ScalableTanh(4)]) for _ in range(4)]
) for n in range(2*2*2)]
length = 4
depth = 4
t = flow.TEBD(2,length,layers,depth,p)
bijective(t)
def flowBuilder(n,
numFlow,
innerBuilder=None,
typeLayer=3,
relax=False,
shift=False):
nn = n * 2
op = source.Gaussian([nn]).to(torch.float64)
if innerBuilder is None:
raise Exception("innerBuilder is None")
if relax:
f3 = flow.DiagScaling(nn,
initValue=0.1 * np.random.randn(nn),
fix=[0] * n + [0] * n,
shift=shift)
else:
f3 = flow.DiagScaling(nn,
initValue=0.1 * np.random.randn(nn),
fix=[0] * n + [1] * n,
shift=shift)
layers = [f3]
if typeLayer == 0:
layers.append(flow.Symplectic(nn))
else:
for d in range(numFlow):
if typeLayer == 3:
layers.append(flow.PointTransformation(innerBuilder(n)))
layers.append(flow.Symplectic(nn))
elif typeLayer == 2:
layers.append(flow.Symplectic(nn))
elif typeLayer == 1:
layers.append(flow.PointTransformation(innerBuilder(n)))
elif typeLayer != 0:
raise Exception("No such type")
return flow.FlowNet(layers, op).double()
def test_bijective():
p = source.Gaussian([4])
f = flow.DiagScaling(4, prior=p)
bijective(f)
def test_saveload():
p = source.Gaussian([1, 4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
repeat = 2
t = flow.MERA(2, length, layers, repeat, prior=p)
p = source.Gaussian([1, 4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
repeat = 2
blankt = flow.MERA(2, length, layers, repeat, prior=p)
saveload(t, blankt)
def test_bijective():
p = source.Gaussian([6])
f = flow.Symplectic(6, prior=p)
bijective(f)
def test_saveload():
p = source.Gaussian([4])
f = flow.Scaling(4, [2, 3], p)
blankf = flow.Scaling(4, prior=p)
saveload(f, blankf)
def test_bijective():
p = source.Gaussian([4])
f = flow.Scaling(4, [2, 3], p)
bijective(f)
def test_saveload():
p = source.Gaussian([3,2,2])
f = flow.OnebyoneConv(2,2,3,prior=p)
blankf = flow.OnebyoneConv(2,2,3,prior=p)
saveload(f,blankf,decimal=3)
# Put structures
Box1_srt = (round(222 * um / dx), round(0 * um / dy), round(0 * um / dz))
Box1_end = (round(272 * um / dx), round(96 * um / dy), round(96 * um / dz))
#Box = structure.Box(Space, Box1_srt, Box1_end, 4., 1.)
# Set PML and PBC
Space.set_PML({'x': '+-', 'y': '+-', 'z': '+-'}, 10)
#Space.set_PML({'x':'+-','y':'+-','z':'+-'}, 10)
# Save eps, mu and PML data.
#Space.save_PML_parameters('./')
#Space.save_eps_mu(savedir)
# Set the type of input source.
Src = source.Gaussian(dt, wvc, spread, pick_pos, np.float64)
Src.plot_pulse(Tstep, freqs, savedir)
#Src = source.Sine(dt, np.float64)
#wvlen = 200*um
#Src.set_wvlen(wvlen)
# Set source position.
#src_xpos = int(Nx/2)
src_xpos = 40
src_ypos = 40
# Momentum of the source.
k0 = 2 * np.pi / wvlen
phi, theta = 0, np.pi / 6
kx = k0 * np.cos(phi) * np.cos(theta)
ky = k0 * np.cos(phi) * np.sin(theta)
courant = 1. / 4
dt = courant * min(dx, dy, dz) / c
Tsteps = int(sys.argv[1])
wvc = 300 * um
interval = 2
spread = 0.3
pick_pos = 1000
plot_per = 100
wvlens = np.arange(200, 600, interval) * um
freqs = c / wvlens
np.save("./graph/freqs", freqs)
# Set the type of input source.
Src = source.Gaussian(dt, wvc, spread, pick_pos, dtype=np.float32)
#Src.plot_pulse(Tsteps, freqs, savedir)
#Src = source.Sine(dt, np.float64)
#Src.set_wvlen( 20 * um)
#sys.exit()
#------------------------------------------------------------------#
#-------------------------- Call objects --------------------------#
#------------------------------------------------------------------#
TF = space.Basic3D((Nx, Ny, Nz), (dx, dy, dz),
dt,
Tsteps,
np.float32,
np.complex64,
def test_saveload():
p = source.Gaussian([3, 2, 2])
f = flow.ArbitraryRotate(3, prior=p)
blankf = flow.ArbitraryRotate(3, prior=p)
saveload(f, blankf, decimal=3)
def test_bijective():
p = source.Gaussian([3, 2, 2])
f = flow.ArbitraryRotate(3, prior=p)
bijective(f, decimal=3)
def test_symplectic():
p = source.Gaussian([6])
f = flow.Symplectic(6, prior=p)
symplectic(f, decimal=4)
def test_saveload():
p = source.Gaussian([4])
f = flow.DiagScaling(4, initValue=[0.01, 0.02, 0.03, 0.04], prior=p)
blankf = flow.DiagScaling(4, prior=p)
saveload(f, blankf, decimal=5)
def test_bijective():
p = source.Gaussian([3,2,2])
f = flow.OnebyoneConv(2,2,3,prior=p)
bijective(f,decimal=3)
