Tring to generate random vectors that look like eigenvalues of random matrices

main.py

from __future__ import print_function
import argparse
import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image

import numpy
import matplotlib.pyplot as plt

import sinkhorn_pointcloud as spc

# Number of eigenvalues
n = 32

# m n-dimensional (over C) examples
#   2n-dimensional over R

# returns array of shape (m, 2*n)
# each row of form (x1, y1, x2, y2, ..., xn, yn)
# where xi+j*yi are eigenvalues of random matrix
def examples(m, n):
    r = numpy.zeros((m, 2*n), dtype=numpy.float32)
    for i in xrange(m):
        mat = numpy.random.normal(size=(n, n)) + 1j*numpy.random.normal(size=(n, n))
        e = numpy.linalg.eigvals(mat)

        x = numpy.real(e)
        y = numpy.imag(e)

        r[i, ::2] = 0.5+0.5*x
        r[i, 1::2] = 0.5+0.5*y

    return r

# show r[0]
def show(r, i):
    x = r[i, ::2]
    y = r[i, 1::2]
    plt.scatter(x, y)

parser = argparse.ArgumentParser(description='VAE Random Eigenvalue Example')
parser.add_argument('--batch-size', type=int, default=128, metavar='N',
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                    help='how many batches to wait before logging training status')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()


torch.manual_seed(args.seed)

device = torch.device("cuda" if args.cuda else "cpu")

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(2*n, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 2*n)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = torch.exp(0.5*logvar)
            eps = torch.randn_like(std)
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return F.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x.view(-1, 2*n))
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar


model = VAE().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function(recon_x, x, mu, logvar):
    epsilon = 0.01
    niter = 100

    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    # I think I need to sum over all rows
    # Also, shouldn't these be arranged as sets of 2D points, not just 2n-d arrays?
    W = 0.0
    for i in xrange(x.shape[0]):
        W += spc.sinkhorn_loss(recon_x[i], x[i].view(-1, 2*n), epsilon, 2*n, niter)

    # How should I weight between KLD and W. Surely not 1:1.
    return KLD+W


def train(epoch, batch_size):
    model.train()
    train_loss = 0
    for batch_idx in xrange(100):
        data = torch.from_numpy(examples(batch_size, n))
        data = data.to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), 0, 
                0,
                loss.item() / len(data)))

    print('====> Epoch: {} loss: {:.4f}'.format(
          epoch, train_loss))


# No need for separate test set. We can generate as many new samples as we like.

for epoch in range(1, args.epochs + 1):
    train(epoch, args.batch_size)

with torch.no_grad():
    sample = torch.randn(64, 20).to(device)
    sample = model.decode(sample).cpu()

for i in xrange(20):
    show(sample, i)
plt.show()