text:
# -*- coding: utf-8 -*-
"""DD-pseudolabels.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1eYKQIIS4i7RPTeAFX5SgIZmG3S2jtnr3

## Double Descent Risk Curve for Neural Networks
The goal of this experiment is to analyze error/loss curve for Resnet18 trained using pseudo-labeling method with CIFAR-10.

- models: resnet18
- dataset: CIFAR-10
- learning algorithm: standard supervised, semi-supervised(pseudo-labeling)
- output : model complexity (number of parameters, epochs) vs. test error/loss
---
- hypothesis: in both supervised and semi-supervised, we should get Double Descent phenomenon, which is defined by
    - having a U-shaped curve before the interpolation threshold (under-parameterized)
    - peaking at the threshold
    - decreasing again in the over-parameterized regime
"""

# Commented out IPython magic to ensure Python compatibility.
from google.colab import drive
drive.mount('/gdrive')

# %matplotlib inline
import torch
from torch.utils.data import DataLoader, Subset, random_split
from torch import nn
import torch.nn.functional as F

from torchvision import datasets, transforms
from torchvision.models.resnet import ResNet, BasicBlock
from torchvision.datasets import CIFAR10, MNIST
from torchvision.transforms import ToTensor

import numpy as np
import pickle, time, json
import matplotlib.pyplot as plt
from PIL import Image

torch.manual_seed(0)
np.random.seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

"""Define some functions"""

def alpha_weight(epoch, T1, T2, af):
    """ calculate value of alpha used in loss based on the epoch

    params:
        - epoch: your current epoch
        - T1: threshold for training with only labeled data
        - T2: threshold for training with only unlabeled data
        - af: max alpha value
    """

    if epoch < T1:
        return 0.0

    elif epoch > T2:
        return af

    else:
        return ((epoch - T1) / (T2 - T1)) * af

def evaluate(model, data_loader, b):
    """ evaluate the loss and accuracy of the trained network on test data

    returns:
        - (test_accuracy, test_loss)

    params:
        - model:
        - test_loader:
        - b:
    """
    correct = 0
    total = 0
    running_loss = 0

    model.eval()
    with torch.no_grad():
        for data, labels in data_loader:
            data = data.to(device)
            labels = labels.to(device)

            output = model(data)
            predicted = torch.max(output, 1)[1]
            correct += (predicted == labels).sum()
            total += data.shape[0]

            loss = loss_fn(output, labels, b)
            running_loss += loss.item()

    test_error = 1 - correct/total
    return test_error, running_loss/len(data_loader)

def loss_fn(outputs, labels, b):
    criterion = nn.CrossEntropyLoss()

    return (criterion(outputs, labels) - b).abs() + b

def train_semisuper(epochs, model, optimizer, train_loader, unlabeled_loader, test_loader, b):
    """train model on labeled and unlabeled data using pseudo-labels/self-training

    returns:
        - metrics:

    params:
        - epochs: total epochs for training
        - model:
        - optimizer:
        - train_loader:
        - unlabeled_loader:
        - test_loader:
        - b:
    """
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    metrics = []
    start_ts = time.time()

    T1 = 20 # no contribution of unlabeled data
    T2 = 150 # less contribution of unlabeled data
    af = 3.1 # Instead of using current epoch we use a "step" variable to calculate alpha_weight. This helps the model converge faster

    # step 1: pre-train teacher model for T1 epochs only on labeled data
    print("START SUPERVISED LEARNING")
    train_super(T1, model, optimizer, train_loader, test_loader, b)
    model.train()

    print("START SEMI-SUPERVISED LEARNING")
    # step 2: train model only on unlabeled data for (T2-T1) epochs and both
    for epoch in range(epochs):
        running_loss = 0
        unlabeled_correct = 0
        unlabeled_total = 0
        labeled_correct = 0
        labeled_total = 0

        # generate pseudo labels and train model on unlabeled data
        for (x_unlabeled, y_unlabeled), (x_labeled, y_labeled) in zip(unlabeled_loader, train_loader):

            # genereate the pseudo labels (changed by every weight update) for unlabeled images
            x_unlabeled, y_unlabeled = x_unlabeled.to(device), y_unlabeled.to(device)
            x_labeled, y_labeled = x_labeled.to(device), y_labeled.to(device)

            output_unlabeled = model(x_unlabeled)
            _, pseudo_labels = torch.max(output_unlabeled, 1) # even if the absolute confidence is low, the max class becomes 1

            unlabeled_loss = loss_fn(output_unlabeled, pseudo_labels, b)

            unlabeled_total += x_unlabeled.shape[0]
            unlabeled_correct += (pseudo_labels == y_unlabeled).sum()

            output_labeled = model(x_labeled)
            _, predicted = torch.max(output_labeled, 1)

            labeled_loss = loss_fn(output_labeled, y_labeled, b)

            labeled_total += x_labeled.shape[0]
            labeled_correct += (predicted == y_labeled).sum()

            total_loss = alpha_weight(epoch, T1, T2, af)*unlabeled_loss + labeled_loss

            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

            running_loss += total_loss.item()

        print('pseudo label accuracy={}'.format(unlabeled_correct/unlabeled_total))

        train_loss = running_loss/len(train_loader)

        correct = unlabeled_correct + labeled_correct
        total = unlabeled_total + labeled_total
        train_error = 1 - correct/total

        test_error, test_loss = evaluate(model, test_loader, b)

        metrics.append([train_error.item(), test_error.item(), train_loss, test_loss])
        print('Epoch: {} : Alpha Weights : {:.3f}, Train Error : {:.3f}, Train Loss : {:.3f} | Test Error : {:.3f}, Test Loss : {:.3f} '.format(epoch+1, alpha_weight(epoch, T1, T2, af), train_error, train_loss, test_error, test_loss))
        model.train()

    print("minutes elapsed: {:.3f}".format((time.time() - start_ts)/60))

    return metrics

def train_super(epochs, model, optimizer, train_loader, test_loader, b):
    """train model in a standard supervised way

    returns:
        - metrics:

    params:
        - epochs: total epochs for training
        - model:
        - optimizer:
        - train_loader:
        - test_loader:
        - b: flooding level
    """
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    metrics = []
    start_ts = time.time()
    model.train()

    # for each epoch
    for epoch in range(epochs):
        running_loss = 0
        total = 0
        correct = 0

        # for each batch
        for batch_idx, (data, labels) in enumerate(train_loader):
            data, labels = data.to(device), labels.to(device)

            output = model(data)
            _, predicted = torch.max(output, 1)

            correct += (predicted == labels).sum()
            total += data.shape[0]

            loss = loss_fn(output, labels, b)
            running_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        train_loss = running_loss/len(train_loader)
        train_error = 1 - correct/total
        test_error, test_loss = evaluate(model, test_loader, b)

        metrics.append([train_error.item(), test_error.item(), train_loss, test_loss])
        print('Epoch: {} : Train Error : {:.3f}, Train Loss : {:.3f} | Test Error : {:.3f}, Test Loss : {:.3f} '.format(epoch+1, train_error, train_loss, test_error, test_loss))
        model.train()

    print("minutes elapsed: {:.3f}".format((time.time() - start_ts)/60))

    return metrics

"""Define your neural net model"""

class PreActBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, **kwargs):
        super(PreActBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)

        if stride != 1 or in_planes != self.expansion * planes:
            self.shortcut = nn.Sequential(nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False))

    def forward(self, x):
        out = F.relu(self.bn1(x))
        shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
        out = self.conv1(out)
        out = self.conv2(F.relu(self.bn2(out)))
        out += shortcut
        return out

class PreActResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10, init_channels=64):
        super(PreActResNet, self).__init__()
        self.in_planes = init_channels
        c = init_channels

        self.conv1 = nn.Conv2d(3, c, kernel_size=3, stride=1, padding=1, bias=False)
        self.layer1 = self._make_layer(block, c, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 2*c, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 4*c, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 8*c, num_blocks[3], stride=2)
        self.dropout = nn.Dropout(0.5)
        self.linear = nn.Linear(8 * c * block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        # eg: [2, 1, 1, ..., 1]. Only the first one downsamples.
        strides = [stride] + [1] * (num_blocks-1)
        layers = []

        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion

        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.dropout(out)
        out = self.linear(out)
        return out

class MNISTResNet(PreActResNet):
    def __init__(self, n_classes, k):
        super(MNISTResNet, self).__init__(PreActBlock, [2, 2, 2, 2], num_classes=n_classes, init_channels=k)

class CIFARResNet(PreActResNet):
    def __init__(self, n_classes, k):
        super(CIFARResNet, self).__init__(PreActBlock, [2, 2, 2, 2], num_classes=n_classes, init_channels=k)

"""#Supervised Experiments

### Experiment 1 - Model-wise DD (Finished!)
"""

basic_setting = {
        'k': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 40, 48, 56, 64],
        'epochs': 300,
        'label_noise': 0.15,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.15,
        'n_labeled': (10000, 40000),
        'augmentation': True
        }

super_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Supervised Experiments/super_model_64.json'
open_file = open(file_name, "w")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

# load either MNIST or CIFAR-10
train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
total_samples = len(train.targets)

# assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
rands = np.random.choice(total_samples, int(basic_setting['label_noise']*total_samples), replace=False)
for rand in rands:
    train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

n_labeled, n_unlabeled = basic_setting['n_labeled']
# split training data into labeled and unlabeled
train, val = random_split(train, [n_labeled, n_unlabeled])
print(len(train), len(val))

train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

print(basic_setting)
for k in basic_setting['k']:
    model = CIFARResNet(basic_setting['n_classes'], k) # define model with the number of parameter
    model.to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print("number of model parameters = {} when k={}".format(total_params, k))

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, basic_setting['b'])
    super_results[str(k)] = error_metrics

    # semi-supervised training using pseudo-labels
    # error_metrics = train_semisuper(basic_setting['epochs'], model, optimizer, train_loader, unlabeled_loader, test_loader, basic_setting['b'])

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""### Experiment 2 - Epoch-wise DD vs. n_samples"""

basic_setting = {
        'k': 64,
        'epochs': 300,
        'label_noise': 0.2,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.15,
        'n_labeled': [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
        'augmentation': True
        }

super_results = {}
semisuper_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/super_epoch_n_samples.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

print(basic_setting)
n_labeled, n_unlabeled = basic_setting['n_labeled']

for ratio in basic_setting['n_labeled']:
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    total_samples = len(train)

    # assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
    rands = np.random.choice(total_samples, int(basic_setting['label_noise']*total_samples), replace=False)
    for rand in rands:
        train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

    # split training data into labeled and unlabeled
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, basic_setting['b'])
    super_results[str(ratio)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""### Experiment 3 - Epoch-wise DD vs. flooding (Finished!)"""

basic_setting = {
        'k': 64,
        'epochs': 200,
        'label_noise': 0.2,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': [0.1, 0.15, 0.2],
        'n_labeled': (50000, 0),
        'augmentation': True
        }

super_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Supervised Experiments/super_epoch_flooding.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

print(basic_setting)

train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    total_samples = len(train)

    # assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
    rands = np.random.choice(total_samples, int(basic_setting['label_noise']*total_samples), replace=False)

    for rand in rands:
        train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

    # split training data into labeled and unlabeled
    n_labeled, n_unlabeled = basic_setting['n_labeled']
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

for flood in basic_setting['b']:

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, flood)
    super_results[str(flood)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""### Experiment 4 - Epoch-wise DD vs. label noise (Finished!)"""

basic_setting = {
        'k': 64,
        'epochs': 200,
        'noise': [0.1, 0.15, 0.2],
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.1,
        'n_labeled': (20000, 30000),
        'augmentation': True
        }

super_results = {}
semisuper_results = {}

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Supervised Experiments/super_epoch_label_noise.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

print(basic_setting)
n_labeled, n_unlabeled = basic_setting['n_labeled']

for noise in basic_setting['noise']:
    # train = datasets.MNIST(root='./data', train=True, download=True, transform=transform_cifar)
    # test = datasets.MNIST(root='./data', train=False, download=True, transform=transform_test)
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    total_samples = len(train)

    # assign random labels to (label_noise)% of the training set (needed for semi-supervised learning)
    rands = np.random.choice(total_samples, int(noise*total_samples), replace=False)

    for rand in rands:
        train.targets[rand] = torch.randint(high=10, size=(1,1)).item()

    # split training data into labeled and unlabeled
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # standard supervised training
    error_metrics = train_super(basic_setting['epochs'], model, optimizer, train_loader, test_loader, basic_setting['b'])
    super_results[str(noise)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(super_results, f)

"""# Semi-Supervised Experiments

### Experiment 1 - Epoch-wise DD vs. labeled ratio
"""

basic_setting = {
        'k': 64,
        'epochs': 200,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.1,
        'n_labeled': [(20000, 30000),(10000, 40000)],
        'augmentation': True
        }

# We observe all forms of double descent most strongly in settings with label noise in the train set (as is often the case when collecting train data in the real-world).
# lr = (n_unlabeled)**(-0.5) # SGD learning rate

print(basic_setting)
semisuper_results = {}
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Semi-supervised Experiments/semisuper_epoch_ratio.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

# load either MNIST or CIFAR-10
# train = datasets.MNIST(root='./data', train=True, download=True, transform=transform_mnist)
# test = datasets.MNIST(root='./data', train=False, download=True, transform=transform_mnist)

for n_labeled, n_unlabeled in basic_setting['n_labeled']:
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=int(len(train)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    unlabeled_loader = DataLoader(val, batch_size=int(len(val)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    # total_params = sum(p.numel() for p in model.parameters())
    # print("number of model parameters = {} when k={}".format(total_params, basic_setting['k']))

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    # semi-supervised training using pseudo-labels
    error_metrics = train_semisuper(basic_setting['epochs'], model, optimizer, train_loader, unlabeled_loader, test_loader, basic_setting['b'])

semisuper_results = {}
semisuper_results[str(n_labeled)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(semisuper_results, f)

"""### Experiment 2 - Model-wise DD"""

basic_setting = {
        'k': 64,
        'epochs': 300,
        'n_batch': 128,
        'n_classes': 10,
        'lr': 1e-4,
        'b': 0.15,
        'n_labeled': [0.6, 0.4, 0.2],
        'augmentation': True
        }
# We observe all forms of double descent most strongly in settings with label noise in the train set (as is often the case when collecting train data in the real-world).
# lr = (n_unlabeled)**(-0.5) # SGD learning rate

print(basic_setting)
semisuper_results = {}
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/semisuper_epoch.json'
open_file = open(file_name, "ab")

# define transformations for training and test set
if basic_setting['augmentation']:
    transform_cifar = transforms.Compose([transforms.ToTensor(),transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
else:
    transform_cifar = transforms.Compose([transforms.ToTensor()])
transform_test = transforms.Compose([transforms.ToTensor()])

# load either MNIST or CIFAR-10
# train = datasets.MNIST(root='./data', train=True, download=True, transform=transform_mnist)
# test = datasets.MNIST(root='./data', train=False, download=True, transform=transform_mnist)

for ratio in basic_setting['n_labeled']:
    train = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_cifar)
    test = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

    # split training data into labeled and unlabeled
    n_labeled = int(total_samples*ratio)
    n_unlabeled = int(total_samples*(1-ratio))
    train, val = random_split(train, [n_labeled, n_unlabeled])
    print("number of labeled: {}, number of unlabeled: {}\n".format(len(train), len(val)))

    train_loader = DataLoader(train, batch_size=int(len(train)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    unlabeled_loader = DataLoader(val, batch_size=int(len(val)/basic_setting['n_batch']), shuffle=True, num_workers=2)
    test_loader = DataLoader(test, batch_size=basic_setting['n_batch'], shuffle=True, num_workers=2)

    model = CIFARResNet(basic_setting['n_classes'], basic_setting['k']) # define model with the number of parameter
    model.to(device)

    # total_params = sum(p.numel() for p in model.parameters())
    # print("number of model parameters = {} when k={}".format(total_params, basic_setting['k']))

    optimizer = torch.optim.Adam(model.parameters(), lr=basic_setting['lr'])
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    # semi-supervised training using pseudo-labels
    error_metrics = train_semisuper(basic_setting['epochs'], model, optimizer, train_loader, unlabeled_loader, test_loader, basic_setting['b'])
    semi_super_results[str(ratio)] = error_metrics

# save list to pickle file
with open(file_name, 'w') as f:
    json.dump(semi_super_results, f)

"""# Plotting"""

# matplotlib.rcParams.update({'font.size': 25})

def plot_modelwise(results, fname):
    titles = ['error', 'loss']
    colors = ['blue', 'lime']
    labels = ['train', 'test']

    k = [] # fixed x-axis
    metrics =. []
    exp_type = fname.split('/')[-2]
    exp_name = fname.split('/')[-1].split('.')[0]

    # extract model size and metrics into separate lists
    for model_size, metric in results.items():
        k.append(int(model_size))
        metrics.append(metric[-1])

    metrics = np.array(metrics)

    # test/train error and loss
    for i, title in enumerate(titles):
        fig, axes = plt.subplots(figsize=(8, 6), dpi=300)
        axes.grid()

        axes.plot(k, metrics[:, 2*i], label=labels[0], color=colors[0]) # train
        axes.plot(k, metrics[:, 2*i+1], label=labels[1], color=colors[1]) # test

        axes.set_xlabel('resnet width=k')
        axes.set_ylabel(title)
        axes.set_title("model width vs. {}".format(title))
        axes.legend(loc='upper right', prop={'size': 15})
        fig.savefig('/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/{}/{}_{}.png'.format(exp_type, exp_name, title), dpi=300)

def plot_epochwise(results, fname):
    titles = ['error', 'loss']
    colors = ['red', 'lime', 'blue']
    labels = ['train', 'test']

    exp_type = fname.split('/')[-2]
    exp_name = fname.split('/')[-1].split('.')[0]

    for key in results.keys():
        epochs = np.arange(1, len(results[key])+1) # fixed x-axis

    for i, title in enumerate(titles):
        fig, axes = plt.subplots(figsize=(8, 6))

        for j, (param, metrics) in enumerate(results.items()):
            metrics = np.array(results[param])
            axes.plot(epochs, metrics[:, ], label=param, color=colors[0])
            axes.plot(epochs, metrics[:, ], label=param, color=colors[1])

        axes.set_xlabel('Epochs')
        axes.set_ylabel(title)
        axes.set_title('Epochs vs. {}'.format(title))
        axes.legend(loc='upper right', prop={'size': 15})
        fig.savefig('/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/{}/{}_{}.png'.format(exp_type, exp_name, title), dpi=300)

file_name = '/gdrive/My Drive/CMSC 828W Research/Code (Won & Amartya)/Semi-supervised Experiments/semisuper_epoch_ratio.json'

with open(file_name) as json_file:
    data = json.load(json_file)
    print(data)
    #plot_modelwise(data, file_name)
    plot_epochwise(data, file_name)

max_stars_repo_name: Ophir-Gal/ssl-descent

max_issues_repo_name: Ophir-Gal/ssl-descent

max_forks_repo_name: Ophir-Gal/ssl-descent


text:
import chainer
import numpy
import pytest
import torch

import chainer_pytorch_migration as cpm

@pytest.mark.parametrize('shape', [(3, 2), (2, 0, 1)])
def test_chainer_parameter(shape):
    # initialized parameter
    arr = numpy.full(shape, 17, 'float32')
    chainer_param = chainer.Parameter(arr)

    # Conversion
    torch_param = cpm.ChainerParameter(chainer_param)

    assert isinstance(torch_param, torch.nn.Parameter)
    assert torch_param.shape == shape
    assert (torch_param.data.numpy() == numpy.full(shape, 17, 'float32')).all()

    # Test memory sharing
    new_arr = numpy.random.randint(-4, 4, shape)
    torch_param.data[...] = torch.tensor(new_arr.copy())
    assert (chainer_param.array == new_arr).all()

def test_chainer_parameter_uninitialized():
    # Uninitialized parameters are not supported
    chainer_param = chainer.Parameter()

    with pytest.raises(TypeError):
        cpm.ChainerParameter(chainer_param)

@pytest.mark.parametrize('shape', [(3, 2), (2, 0, 1)])
def test_chainer_parameter_grad_getter(shape):
    arr = numpy.full(shape, 17, 'float32')
    grad = numpy.full(shape, 9, 'float32')
    chainer_param = chainer.Parameter(arr)
    chainer_param.grad = grad.copy()

    # Conversion
    torch_param = cpm.ChainerParameter(chainer_param)

    # Getter
    torch_grad = torch_param.grad

    assert isinstance(torch_grad, torch.Tensor)
    assert (torch_grad.numpy() == grad).all()

    # Test memory sharing
    new_arr = numpy.random.randint(-4, 4, shape)
    torch_grad[...] = torch.tensor(new_arr.copy())
    assert (chainer_param.grad == new_arr).all()

@pytest.mark.parametrize('shape', [(3, 2), (2, 0, 1)])
def test_chainer_parameter_grad_setter(shape):
    arr = numpy.full(shape, 17, 'float32')
    chainer_param = chainer.Parameter(arr)

    # Conversion
    torch_param = cpm.ChainerParameter(chainer_param)
    # Initialize grad
    torch_param.requires_grad = True
    optimizer = torch.optim.SGD([torch_param], lr=0.01, momentum=0.9)
    optimizer.zero_grad()

    # Setter
    grad = torch.full(shape, 9, dtype=torch.float32)
    torch_param.grad = grad
    numpy.testing.assert_array_equal(grad, torch_param.grad)

def test_to_chainer_parameters():
    # initialized parameter
    a_arr = numpy.ones((3,  2), 'float32')
    a_chainer_param = chainer.Parameter(a_arr)
    # 0-size parameter
    b_arr = numpy.ones((2, 0, 1), 'float32')
    b_chainer_param = chainer.Parameter(b_arr)

    link = chainer.Link()
    with link.init_scope():
        link.a = a_chainer_param
        link.b = b_chainer_param

    # Conversion
    torched = cpm.LinkAsTorchModel(link)
    params = list(torched.parameters())
    assert len(params) == 2
    assert isinstance(params[0], torch.nn.Parameter)
    assert isinstance(params[1], torch.nn.Parameter)
    assert params[0].shape == (3, 2)
    assert params[1].shape == (2, 0, 1)
    assert (params[0].data.numpy() == numpy.ones((3, 2))).all()

    # Test memory sharing
    params[0].data[...] = torch.tensor(numpy.arange(6).reshape((3, 2)))
    assert (a_chainer_param.array == numpy.arange(6).reshape((3, 2))).all()

def test_to_chainer_parameters_uninitialized():
    # Uninitialized parameters are not supported
    a_chainer_param = chainer.Parameter()

    link = chainer.Link()
    with link.init_scope():
        link.a = a_chainer_param

    with pytest.raises(TypeError):
        torched = cpm.LinkAsTorchModel(link)
        torched.parameters()

def test_state_dict():
    a_arr = numpy.ones((3,  2), 'float32')
    a_chainer_param = chainer.Parameter(a_arr)
    # 0-size parameter
    b_arr = numpy.ones((2, 0, 1), 'float32')
    b_chainer_param = chainer.Parameter(b_arr)

    link = chainer.Link()
    with link.init_scope():
        link.a = a_chainer_param
        link.b = b_chainer_param

    torched = cpm.LinkAsTorchModel(link)
    state_dict = torched.state_dict()
    assert '/a' in state_dict
    numpy.testing.assert_array_equal(a_arr, state_dict['/a'].detach())
    assert '/b' in state_dict
    numpy.testing.assert_array_equal(b_arr, state_dict['/b'].detach())

def test_named_params():
    a_arr = numpy.ones((3,  2), 'float32')
    a_chainer_param = chainer.Parameter(a_arr)
    # 0-size parameter
    b_arr = numpy.ones((2, 0, 1), 'float32')
    b_chainer_param = chainer.Parameter(b_arr)

    link = chainer.Link()
    with link.init_scope():
        link.a = a_chainer_param
        link.b = b_chainer_param

    torched = cpm.LinkAsTorchModel(link)
    n_params = dict(torched.named_parameters())
    assert '/a' in n_params
    numpy.testing.assert_array_equal(a_arr, n_params['/a'].detach())
    assert '/b' in n_params
    numpy.testing.assert_array_equal(b_arr, n_params['/b'].detach())

ext: py

max_stars_repo_name: emcastillo/chainer-pytorch-migration

max_issues_repo_name: emcastillo/chainer-pytorch-migration

max_forks_repo_name: emcastillo/chainer-pytorch-migration

reason: import numpy


text:
"""
.. module:: reporters
   :platform: Unix, Windows
   :synopsis: a module for defining OpenMM reporter classes.

.. moduleauthor:: Charlles R. A. Abreu <abreu@eq.ufrj.br>

.. _pandas.DataFrame: https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.html
.. _StateDataReporter: http://docs.openmm.org/latest/api-python/generated/simtk.openmm.app.statedatareporter.StateDataReporter.html
.. _CustomIntegrator: http://docs.openmm.org/latest/api-python/generated/simtk.openmm.openmm.CustomIntegrator.html
.. _CustomCVForce: docs.openmm.org/latest/api-python/generated/simtk.openmm.openmm.CustomCVForce.html

"""

import sys

import numpy as np
import pandas as pd
from simtk import openmm
from simtk import unit
from simtk.openmm import app

from .computers import PressureComputer
from .computers import _MoleculeTotalizer
from .utils import InputError

class _MultiStream:
    def __init__(self, outputs):
        self._outputs = list()
        for output in outputs:
            self._outputs.append(open(output, 'w') if isinstance(output, str) else output)

    def __del__(self):
        for output in self._outputs:
            if output != sys.stdout and output != sys.stderr:
                output.close()

    def write(self, message):
        for output in self._outputs:
            output.write(message)

    def flush(self):
        for output in self._outputs:
            output.flush()

class _AtomsMM_Reporter():
    """
    Base class for reporters.

    """
    def __init__(self, file, reportInterval, **kwargs):
        self._reportInterval = reportInterval
        self._requiresInitialization = True
        self._needsPositions = False
        self._needsVelocities = False
        self._needsForces = False
        self._needEnergy = False
        extraFile = kwargs.pop('extraFile', None)
        if extraFile is None:
            self._out = open(file, 'w') if isinstance(file, str) else file
        else:
            self._out = _MultiStream([file, extraFile])
        self._separator = kwargs.pop('separator', ',')

    def _initialize(self, simulation, state):
        pass

    def _generateReport(self, simulation, state):
        pass

    def describeNextReport(self, simulation):
        """
        Get information about the next report this object will generate.

        Parameters
        ----------
            simulation : Simulation
                The Simulation to generate a report for

        Returns
        -------
            tuple
                A five element tuple. The first element is the number of steps
                until the next report. The remaining elements specify whether
                that report will require positions, velocities, forces, and
                energies respectively.

        """
        steps = self._reportInterval - simulation.currentStep % self._reportInterval
        return (steps, self._needsPositions, self._needsVelocities, self._needsForces, self._needEnergy)

    def report(self, simulation, state):
        """
        Generate a report.

        Parameters
        ----------
            simulation : Simulation
                The Simulation to generate a report for
            state : State
                The current state of the simulation

        """
        if self._requiresInitialization:
            self._initialize(simulation, state)
            self._requiresInitialization = False
        self._generateReport(simulation, state)

class ExtendedStateDataReporter(app.StateDataReporter):
    """
    An extension of OpenMM's StateDataReporter_ class, which outputs information about a simulation,
    such as energy and temperature, to a file.

    All original functionalities of StateDataReporter_ are preserved and the following ones are
    included:

    1. Report the Coulomb contribution of the potential energy (keyword: `coulombEnergy`):

        This contribution includes both real- and reciprocal-space terms.

    2. Report the atomic virial of a fully-flexible system (keyword: `atomicVirial`):

        Considering full scaling of atomic coordinates in a box volume change (i.e. without any
        distance constraints), the internal virial of the system is given by

        .. math::
            W = -\\sum_{i,j} r_{ij} E^\\prime(r_{ij}),

        where :math:`E^\\prime(r)` is the derivative of the pairwise interaction potential as a
        function of the distance between to atoms. Such interaction includes van der Waals, Coulomb,
        and bond-stretching contributions. Bond-bending and dihedral angles are not considered
        because they are invariant to full volume-scaling of atomic coordinates.

    3. Report the nonbonded contribution of the atomic virial (keyword: `nonbondedVirial`):

        The nonbonded virial is given by

        .. math::
            W_\\mathrm{nb} = -\\sum_{i,j} r_{ij} E_\\mathrm{nb}^\\prime(r_{ij}),

        where :math:`E_\\mathrm{nb}^\\prime(r)` is the derivative of the nonbonded pairwise
        potential, which comprises van der Waals and Coulomb interactions only.

    4. Report the atomic pressure of a fully-flexible system (keyword: `atomicPressure`):

        .. math::
            P = \\frac{2 K + W}{3 V},

        where :math:`K` is the kinetic energy sum for all atoms in the system. If keyword
        `bathTemperature` is employed (see below), the instantaneous kinetic energy is substituted
        by its equipartition-theorem average
        :math:`\\left\\langle K \\right\\rangle = 3 N_\\mathrm{atoms} k_B T/2`,
        where :math:`T` is the heat-bath temperature.

    5. Report the molecular virial of a system (keyword: `molecularVirial`):

        To compute the molecular virial, only the center-of-mass coordinates of the molecules are
        considered to scale in a box volume change, while the internal molecular structure is kept
        unaltered. The molecular virial is computed from the nonbonded part of the atomic virial by
        using the formulation of Ref. :cite:`Hunenberger_2002`:

        .. math::
            W_\\mathrm{mol} = W - \\sum_{i} (\\mathbf{r}_i - \\mathbf{r}_i^\\mathrm{cm}) \\cdot \\mathbf{F}_i,

        where :math:`\\mathbf{r}_i` is the coordinate of atom i, :math:`\\mathbf{F}_i` is the
        resultant pairwise force acting on it (excluding bond-bending and dihedral angles), and
        :math:`\\mathbf{r}_i^\\mathrm{cm}` is the center-of-mass coordinate of its containing
        molecule.

    6. Report the molecular pressure of a system (keyword: `molecularPressure`):

        .. math::
            P = \\frac{2 K_\\mathrm{mol} + W_\\mathrm{mol}}{3 V},

        where :math:`K_\\mathrm{mol}` is the center-of-mass kinetic energy summed for all molecules
        in the system. If keyword `bathTemperature` is employed (see below), the instantaneous
        kinetic energy is substituted by its equipartition-theorem average
        :math:`\\left\\langle K_\\mathrm{mol} \\right\\rangle = 3 N_\\mathrm{mols} k_B T/2`,
        where :math:`T` is the heat-bath temperature.

    7. Report the center-of-mass kinetic energy (keyword: `molecularKineticEnergy`):

        .. math::
            K_\\mathrm{mol} = \\frac{1}{2} \\sum_{i=1}^{N_\\mathrm{mol}} M_i v_{\\mathrm{cm}, i}^2,

        where :math:`N_\\mathrm{mol}` is the number of molecules in the system, :math:`M_i` is the
        total mass of molecule `i`, and :math:`v_{\\mathrm{cm}, i}` is the center-of-mass velocity
        of molecule `i`.

    8. Report potential energies at multiple global parameter states (keyword: `globalParameterStates`):

        Computes and reports the potential energy of the system at a number of provided global
        parameter states.

    9. Report global parameter values (keyword: `globalParameters`):

        Reports the values of specified global parameters.

    10. Report derivatives of energy with respect to global parameters (keyword: `energyDerivatives`):

        Computes and reports derivatives of the potential energy of the system at the current
        state with respect to specified global parameters.

    11. Report values of collective variables (keyword: `collectiveVariables`)

        Report the values of a set of collective variables.

    12. Allow specification of an extra file for reporting (keyword: `extraFile`).

        This can be used for replicating a report simultaneously to `sys.stdout` and to a file
        using a unique reporter.

    Keyword Args
    ------------
        coulombEnergy : bool, optional, default=False
            Whether to write the Coulomb contribution of the potential energy to the file.
        atomicVirial : bool, optional, default=False
            Whether to write the total atomic virial to the file.
        nonbondedVirial : bool, optional, default=False
            Whether to write the nonbonded contribution to the atomic virial to the file.
        atomicPressure : bool, optional, default=False
            Whether to write the internal atomic pressure to the file.
        molecularVirial : bool, optional, default=False
            Whether to write the molecular virial to the file.
        molecularPressure : bool, optional, default=False
            Whether to write the internal molecular pressure to the file.
        molecularKineticEnergy : bool, optional, default=False
            Whether to write the molecular center-of-mass kinetic energy to the file.
        globalParameterStates : pandas.DataFrame_, optional, default=None
            A DataFrame containing context global parameters (column names) and sets of values
            thereof. If it is provided, then the potential energy will be reported for every state
            these parameters define.
        globalParameters : list(str), optional, default=None
            A list of global parameter names. If it is provided, then the values of these parameters
            will be reported.
        energyDerivatives : list(str), optional, default=None
            A list of global parameter names. If it is provided, then the derivatives of the
            total potential energy with respect to these parameters will be reported. It is
            necessary that the calculation of these derivatives has been activated beforehand
            (see, for instance, CustomIntegrator_).
        collectiveVariables : list(openmm.CustomCVForce), optional, default=None
            A list of CustomCVForce_ objects. If it is provided, then the values of all collective
            variables associated with these objects will be reported.
        pressureComputer : :class:`~atomsmm.computers.PressureComputer`, optional, default=None
            A computer designed to determine pressures and virials. This is mandatory if any keyword
            related to virial or pressure is set as `True`.
        extraFile : str or file, optional, default=None
            Extra file to write to, specified as a file name or a file object.

    """
    def __init__(self, file, reportInterval, **kwargs):
        self._coulombEnergy = kwargs.pop('coulombEnergy', False)
        self._atomicVirial = kwargs.pop('atomicVirial', False)
        self._nonbondedVirial = kwargs.pop('nonbondedVirial', False)
        self._atomicPressure = kwargs.pop('atomicPressure', False)
        self._molecularVirial = kwargs.pop('molecularVirial', False)
        self._molecularPressure = kwargs.pop('molecularPressure', False)
        self._molecularKineticEnergy = kwargs.pop('molecularKineticEnergy', False)
        self._globalParameterStates = kwargs.pop('globalParameterStates', None)
        self._globalParameters = kwargs.pop('globalParameters', None)
        self._energyDerivatives = kwargs.pop('energyDerivatives', None)
        self._collectiveVariables = kwargs.pop('collectiveVariables', None)
        self._pressureComputer = kwargs.pop('pressureComputer', None)
        extra = kwargs.pop('extraFile', None)
        if extra is None:
            super().__init__(file, reportInterval, **kwargs)
        else:
            super().__init__(_MultiStream([file, extra]), reportInterval, **kwargs)
        self._computing = any([self._coulombEnergy,
                               self._atomicVirial,
                               self._nonbondedVirial,
                               self._atomicPressure,
                               self._molecularVirial,
                               self._molecularPressure,
                               self._molecularKineticEnergy])
        if self._computing:
            if self._pressureComputer is not None and not isinstance(self._pressureComputer, PressureComputer):
                raise InputError('keyword "pressureComputer" requires a PressureComputer instance')
            self._needsPositions = True
            self._needsForces = any([self._needsForces,
                                     self._molecularVirial,
                                     self._molecularPressure])
            self._needsVelocities = any([self._needsVelocities,
                                         self._molecularPressure,
                                         self._atomicPressure,
                                         self._molecularKineticEnergy])
        self._backSteps = -sum([self._speed, self._elapsedTime, self._remainingTime])

    def _add_item(self, lst, item):
        if self._backSteps == 0:
            lst.append(item)
        else:
            lst.insert(self._backSteps, item)

    def _constructHeaders(self):
        headers = super()._constructHeaders()
        if self._coulombEnergy:
            self._add_item(headers, 'Coulomb Energy (kJ/mole)')
        if self._atomicVirial:
            self._add_item(headers, 'Atomic Virial (kJ/mole)')
        if self._nonbondedVirial:
            self._add_item(headers, 'Nonbonded Virial (kJ/mole)')
        if self._atomicPressure:
            self._add_item(headers, 'Atomic Pressure (atm)')
        if self._molecularVirial:
            self._add_item(headers, 'Molecular Virial (kJ/mole)')
        if self._molecularPressure:
            self._add_item(headers, 'Molecular Pressure (atm)')
        if self._molecularKineticEnergy:
            self._add_item(headers, 'Molecular Kinetic Energy (kJ/mole)')
        if self._globalParameterStates is not None:
            for index in self._globalParameterStates.index:
                self._add_item(headers, 'Energy[{}] (kJ/mole)'.format(index))
        if self._globalParameters is not None:
            for name in self._globalParameters:
                self._add_item(headers, name)
        if self._energyDerivatives is not None:
            for name in self._energyDerivatives:
                self._add_item(headers, 'diff(E,{})'.format(name))
        if self._collectiveVariables is not None:
            for force in self._collectiveVariables:
                for index in range(force.getNumCollectiveVariables()):
                    name = force.getCollectiveVariableName(index)
                    self._add_item(headers, name)
        return headers

    def _constructReportValues(self, simulation, state):
        values = super()._constructReportValues(simulation, state)
        if self._computing:
            computer = self._pressureComputer
            computer.import_configuration(state)
            atomicVirial = computer.get_atomic_virial().value_in_unit(unit.kilojoules_per_mole)
            if self._coulombEnergy:
                coulombVirial = computer.get_coulomb_virial()
                self._add_item(values, coulombVirial.value_in_unit(unit.kilojoules_per_mole))
            if self._atomicVirial:
                self._add_item(values, atomicVirial)
            if self._nonbondedVirial:
                nonbondedVirial = computer.get_dispersion_virial() + computer.get_coulomb_virial()
                self._add_item(values, nonbondedVirial.value_in_unit(unit.kilojoules_per_mole))
            if self._atomicPressure:
                atomicPressure = computer.get_atomic_pressure()
                self._add_item(values, atomicPressure.value_in_unit(unit.atmospheres))
            if self._molecularVirial or self._molecularPressure:
                forces = state.getForces(asNumpy=True)
                if self._molecularVirial:
                    molecularVirial = computer.get_molecular_virial(forces)
                    self._add_item(values, molecularVirial.value_in_unit(unit.kilojoules_per_mole))
                if self._molecularPressure:
                    molecularPressure = computer.get_molecular_pressure(forces)
                    self._add_item(values, molecularPressure.value_in_unit(unit.atmospheres))
            if self._molecularKineticEnergy:
                molKinEng = computer.get_molecular_kinetic_energy()
                self._add_item(values, molKinEng.value_in_unit(unit.kilojoules_per_mole))

        if self._globalParameterStates is not None:
            original = dict()
            for name in self._globalParameterStates.columns:
                original[name] = simulation.context.getParameter(name)
            latest = original.copy()
            for index, row in self._globalParameterStates.iterrows():
                for name, value in row.items():
                    if value != latest[name]:
                        simulation.context.setParameter(name, value)
                        latest[name] = value
                energy = simulation.context.getState(getEnergy=True).getPotentialEnergy()
                self._add_item(values, energy.value_in_unit(unit.kilojoules_per_mole))
            for name, value in original.items():
                if value != latest[name]:
                    simulation.context.setParameter(name, value)

        if self._globalParameters is not None:
            for name in self._globalParameters:
                self._add_item(values, simulation.context.getParameter(name))

        if self._energyDerivatives is not None:
            mystate = simulation.context.getState(getParameterDerivatives=True)
            derivative = mystate.getEnergyParameterDerivatives()
            for name in self._energyDerivatives:
                self._add_item(values, derivative[name])

        if self._collectiveVariables is not None:
            for force in self._collectiveVariables:
                for cv in force.getCollectiveVariableValues(simulation.context):
                    self._add_item(values, cv)

        return values

class XYZReporter(_AtomsMM_Reporter):
    """
    Outputs to an XYZ-format file a series of frames containing the coordinates, velocities,
    momenta, or forces on all atoms in a Simulation.

    .. note::
        Coordinates are expressed in nanometers, velocities in nanometer/picosecond, momenta in
        dalton*nanometer/picosecond, and forces in dalton*nanometer/picosecond^2.

    To use this reporter, create an XYZReporter object and append it to the Simulation's list of
    reporters.

    Keyword Args
    ------------
        output : str, default='positions'
            Which kind of info to report. Valid options are 'positions', 'velocities', 'momenta' and
            'forces'.
        groups : set(int), default=None
            Which force groups to consider in the force calculations. If this is `None`, then all
            force groups will be evaluated.

    """
    def __init__(self, file, reportInterval, **kwargs):
        self._output = kwargs.get('output', 'positions')
        self._groups = kwargs.get('groups', None)
        if self._output == 'positions':
            self._unit = unit.angstroms
        elif self._output == 'velocities':
            self._unit = unit.angstroms/unit.picoseconds
        elif self._output == 'momenta':
            self._unit = unit.dalton*unit.angstroms/unit.picoseconds
        elif self._output == 'forces':
            self._unit = unit.dalton*unit.angstroms/unit.picoseconds**2
        else:
            raise InputError('Unrecognizable keyword value')
        super().__init__(file, reportInterval, **kwargs)
        self._needsPositions = self._output == 'positions'
        self._needsVelocities = self._output in ['velocities', 'momenta']
        self._needsForces = self._output == 'forces'

    def _initialize(self, simulation, state):
        self._symbols = [atom.element.symbol for atom in simulation.topology.atoms()]
        sys = simulation.system
        self._N = sys.getNumParticles()
        if self._output == 'momenta':
            mass = [sys.getParticleMass(i).value_in_unit(unit.dalton) for i in range(self._N)]
            self._mass = np.vstack([mass, mass, mass]).transpose()*unit.dalton

    def _get_values(self, simulation, state):
        if self._output == 'positions':
            values = state.getPositions(asNumpy=True)
        elif self._output == 'velocities':
            values = state.getVelocities(asNumpy=True)
        elif self._output == 'momenta':
            values = self._mass*state.getVelocities(asNumpy=True)
        elif self._groups is None:
            values = state.getForces(asNumpy=True)
        else:
            new_state = simulation.context.getState(getForces=True, groups=self._groups)
            values = new_state.getForces(asNumpy=True)
        return values.value_in_unit(self._unit)

    def _write(self, step, N, names, values):
        print(N, file=self._out)
        pd.DataFrame(index=names, data=values).to_csv(
            self._out,
            sep='\t',
            header=[f'{self._output} in {self._unit} at time step {step}', '', ''],
        )

    def _generateReport(self, simulation, state):
        values = self._get_values(simulation, state)
        self._write(simulation.currentStep, self._N, self._symbols, values)

class CenterOfMassReporter(XYZReporter):
    """
    Outputs to an XYZ-format file a series of frames containing the center-of-mass coordinates,
    center-of-mass velocities, total momenta, or resultant forces on all molecules in a Simulation.

    .. note::
        Coordinates are expressed in nanometers, velocities in nanometer/picosecond, momenta in
        dalton*nanometer/picosecond, and forces in dalton*nanometer/picosecond^2.

    To use this reporter, create an CenterOfMassReporter object and append it to the Simulation's
    list of reporters.

    Keyword Args
    ------------
        output : str, default='positions'
            Which kind of info to report. Valid options are 'positions', 'velocities', 'momenta' and
            'forces'.
        groups : set(int), default=None
            Which force groups to consider in the force calculations. If this is `None`, then all
            force groups will be evaluated.

    """
    def _initialize(self, simulation, state):
        super()._initialize(simulation, state)
        self._mols = _MoleculeTotalizer(simulation.context, simulation.topology)

    def _generateReport(self, simulation, state):
        values = self._get_values(simulation, state)
        if self._output in ['positions', 'velocities']:
            cm_values = self._mols.massFrac.dot(values)
        else:
            cm_values = self._mols.selection.dot(values)
        self._write(simulation.currentStep, self._mols.nmols, self._mols.residues, cm_values)

class CustomIntegratorReporter(_AtomsMM_Reporter):
    """
    Outputs global and per-DoF variables of a CustomIntegrator instance.

    Keyword Args
    ------------
        describeOnly : bool, optional, default=True
            Whether to output only descriptive statistics that summarize the activated per-Dof
            variables.

    """
    def __init__(self, file, reportInterval, **kwargs):
        super().__init__(file, reportInterval, **kwargs)
        self._describeOnly = kwargs.pop('describeOnly', True)
        self._variables = []
        for key, value in kwargs.items():
            if value is True:
                self._variables.append(key)
        if not self._variables:
            raise InputError("No global or perDof variables have been passed")

    def _initialize(self, simulation, state):
        integrator = self._integrator = simulation.integrator
        if not isinstance(integrator, openmm.CustomIntegrator):
            raise Exception("simulation.integrator is not a CustomIntegrator")
        self._globals = {}
        for index in range(integrator.getNumGlobalVariables()):
            variable = integrator.getGlobalVariableName(index)
            if variable in self._variables:
                self._globals[variable] = index
        self._perDof = {}
        for index in range(integrator.getNumPerDofVariables()):
            variable = integrator.getPerDofVariableName(index)
            if variable in self._variables:
                self._perDof[variable] = index
        if set(self._variables) != set(self._globals) | set(self._perDof):
            raise InputError("Unknown variables have been passed")

    def _generateReport(self, simulation, state):
        for variable, index in self._globals.items():
            value = self._integrator.getGlobalVariable(index)
            print('{}\n{}'.format(variable, value), file=self._out)

        for variable, index in self._perDof.items():
            values = self._integrator.getPerDofVariable(index)
            titles = ['{}.{}'.format(variable, dir) for dir in ['x', 'y', 'z']]
            df = pd.DataFrame(data=np.array(values), columns=titles)
            if self._describeOnly:
                print(df.describe(), file=self._out)
            else:
                df.to_csv(self._out, sep='\t')

class ExpandedEnsembleReporter(_AtomsMM_Reporter):
    """
    Performs an Expanded Ensemble simulation and reports the energies of multiple states.

    Parameters
    ----------
        states : pandas.DataFrame_
            A DataFrame containing context global parameters (column names) and sets of values
            thereof. The potential energy will be reported for every state these parameters define.
            If one of the variables is named as `weight`, then its set of values will be assigned
            to every state as an importance sampling weight. Otherwise, all states will have
            identical weights. States which are supposed to only have their energies reported, with
            no actual visits, can have their weights set up to `-inf`.
        temperature : unit.Quantity
            The system temperature.

    Keyword Args
    ------------
        reportsPerExchange : int, optional, default=1
            The number of reports between attempts to exchange the global parameter state, that is,
            the exchange interval measured in units of report intervals.

    """
    def __init__(self, file, reportInterval, states, temperature, **kwargs):
        self._parameter_states = states.copy()
        self._nstates = len(states.index)
        self._reports_per_exchange = kwargs.pop('reportsPerExchange', 1)
        super().__init__(file, reportInterval, **kwargs)
        if 'weight' in states:
            self._weights = self._parameter_states.pop('weight').values
            finite = np.where(np.isfinite(self._weights))[0]
            self._first_state = finite[0]
            self._last_state = finite[-1]
        else:
            self._weights = np.zeros(self._nstates)
            self._first_state = 0
            self._last_state = self._nstates - 1
        kT = (unit.MOLAR_GAS_CONSTANT_R*temperature).value_in_unit(unit.kilojoules_per_mole)
        self._beta = 1.0/kT
        self._nreports = 0
        self._overall_visits = np.zeros(self._nstates, dtype=int)
        self._downhill_visits = np.zeros(self._nstates, dtype=int)
        self._probability_accumulators = np.zeros(self._nstates)
        self._downhill = False
        self._counting_started = False
        self._regime_change = []

    def _initialize(self, simulation, state):
        headers = ['step', 'state']
        for index in self._parameter_states.index:
            headers.append('Energy[{}] (kJ/mole)'.format(index))
        print(*headers, sep=self._separator, file=self._out)

    def _register_visit(self, state):
        if self._downhill:
            if state == self._first_state:
                self._downhill = False
                self._regime_change.append(self._nreports)
        elif state == self._last_state:
            self._downhill = True
            self._regime_change.append(self._nreports)
        if self._counting_started:
            self._overall_visits[state] += 1
            if self._downhill:
                self._downhill_visits[state] += 1
        else:
            self._counting_started = self._downhill is True

    def _generateReport(self, simulation, state):
        energies = np.zeros(self._nstates)
        original = dict()
        for name in self._parameter_states.columns:
            original[name] = simulation.context.getParameter(name)
        latest = original.copy()
        for i, (index, row) in enumerate(self._parameter_states.iterrows()):
            for name, value in row.items():
                if value != latest[name]:
                    simulation.context.setParameter(name, value)
                    latest[name] = value
            energy = simulation.context.getState(getEnergy=True).getPotentialEnergy()
            energies[i] = energy.value_in_unit(unit.kilojoules_per_mole)
        self._nreports += 1
        exponents = self._weights - self._beta*energies
        probabilities = np.exp(exponents - np.amax(exponents))
        probabilities /= np.sum(probabilities)
        self._probability_accumulators += probabilities
        if self._nreports % self._reports_per_exchange == 0:
            state = np.random.choice(self._nstates, p=probabilities)
            for name, value in self._parameter_states.iloc[state].items():
                if value != latest[name]:
                    simulation.context.setParameter(name, value)
            self._register_visit(state)
        print(simulation.currentStep, state, *energies, sep=self._separator, file=self._out)

    def _isochronal_delta(self, f, n):
        N = len(f)
        b = 3/(n*(n+1)*(2*n+1))
        seq = np.arange(1, n+1)
        a = (b/2)*np.array([n*(n+1)-k*(k-1) for k in seq])
        ind = np.argsort(f)
        fa = f[ind]
        delta = np.empty(N)
        delta[0] = -fa[0]/2 + np.sum(a*fa[1:n+1])
        for i in range(1, N-1):
            delta[i] = b*np.sum([k*(fa[min(i+k, N-1)] - fa[max(i-k, 0)]) for k in seq])
        delta[N-1] = fa[N-1]/2 - np.sum(np.flip(a)*fa[N-n-1:N-1])
        delta[ind] = delta
        return delta

    def read_csv(self, file, **kwargs):
        comment = kwargs.pop('comment', '#')
        separator = kwargs.pop('sep', self._separator)
        df = pd.read_csv(file, comment=comment, sep=separator, **kwargs)
        energies = np.zeros(self._nstates)
        for index, row in df.iterrows():
            state = int(row['state'])
            for i in self._parameter_states.index:
                energies[i] = row['Energy[{}] (kJ/mole)'.format(i)]
            self._nreports += 1
            exponents = self._weights - self._beta*energies
            probabilities = np.exp(exponents - np.amax(exponents))
            probabilities /= np.sum(probabilities)
            self._probability_accumulators += probabilities
            if self._nreports % self._reports_per_exchange == 0:
                self._register_visit(state)

    def state_sampling_analysis(self, staging_variable=None, to_file=True, isochronal_n=2):
        """
        Build histograms of states visited during the overall process as well as during downhill
        walks.

        Returns
        -------
            pandas.DataFrame_

        """
        mask = self._overall_visits > 0
        frame = pd.DataFrame(self._parameter_states)[mask]
        histogram = self._overall_visits[mask]
        downhill_fraction = self._downhill_visits[mask]/histogram
        weight = self._weights[mask]
        frame['weight'] = weight
        frame['histogram'] = histogram/np.sum(histogram)
        frame['downhill_fraction'] = downhill_fraction
        if self._counting_started:
            probability = self._probability_accumulators[mask]/self._nreports
            free_energy = weight - np.log(probability)
            free_energy -= free_energy[0]
            delta = self._isochronal_delta(downhill_fraction, isochronal_n)
            isochronal_weight = weight + 0.5*np.log(delta/probability)
            frame['free_energy'] = free_energy
            frame['isochronal_histogram'] = np.sqrt(delta*probability)
            frame['isochronal_weight'] = isochronal_weight - isochronal_weight[0]
            if staging_variable is not None:
                x = frame[staging_variable].values
                f = downhill_fraction
                n = len(x)
                optimal_pdf = np.sqrt(np.diff(f)/np.diff(x))      # Stepwise optimal PDF
                area = optimal_pdf*np.diff(x)                     # Integral in each interval
                optimal_cdf = np.cumsum(area)/np.sum(area)        # Piecewise linear optimal CDF
                optimal_x = np.interp(np.linspace(0, 1, n), np.insert(optimal_cdf, 0, 0), x)
                frame['staging_{}'.format(staging_variable)] = optimal_x
                frame['staging_weight'] = np.interp(optimal_x, x, free_energy)
        if to_file:
            print('# {0} State Sampling Analysis {0}'.format('-'*40), file=self._out)
            with pd.option_context('display.max_rows', None, 'display.max_columns', None):
                print('# ' + frame.to_string(index=False).replace('\n', '\n# '), file=self._out)
        return frame

    def walking_time_analysis(self, history=False, to_file=True):
        times = np.diff(np.array(self._regime_change))
        downhill = self._reportInterval*times[0::2]
        uphill = self._reportInterval*times[1::2]
        if history:
            df = pd.DataFrame({'downhill': pd.Series(downhill),
                               'uphill': pd.Series(uphill)})
            print('# {0} Walking Time History {0}'.format('-'*10), file=self._out)
            print('# ' + df.to_string().replace('\n', '\n# '), file=self._out)
        df = pd.DataFrame(index=['count', 'mean time'],
                          columns=['downhill', 'uphill'],
                          data=[[downhill.size, uphill.size],
                                [downhill.mean(), uphill.mean()]],
                          dtype='object')
        if to_file:
            print('# {0} Walking Time Analysis {0}'.format('-'*10), file=self._out)
            print('# ' + df.to_string().replace('\n', '\n# '), file=self._out)
        return df

max_stars_repo_name: craabreu/atomsmm

max_issues_repo_name: craabreu/atomsmm

max_forks_repo_name: craabreu/atomsmm


text:
#!/usr/bin/env python

import networkx as nx
import dimod
import matplotlib.pyplot as plt
from dwave.system import LeapHybridDQMSampler
import numpy as np
from random import random

G = nx.karate_club_graph()

dqm = dimod.DiscreteQuadraticModel()
n_clusters = 4
n_nodes = G.number_of_nodes()
beta = 3

for i in G.nodes:
    dqm.add_variable(n_clusters, label=i)

norm_factor = beta * (1 - (2 * n_nodes / n_clusters))
for i in G.nodes:
    linear_term = norm_factor + (0.5 * np.ones(n_clusters) * G.degree[i])
    dqm.set_linear(i, linear_term)

# Quadratic term for node pairs which do not have edges between them
for p0, p1 in nx.non_edges(G):
    dqm.set_quadratic(p0, p1, {(c, c): (2 * beta) for c in range(n_clusters)})

# Quadratic term for node pairs which have edges between them
for p0, p1 in G.edges:
    dqm.set_quadratic(p0, p1, {(c, c): ((2 * beta) - 1) for c in range(n_clusters)})

sampler = LeapHybridDQMSampler()
sampleset = sampler.sample_dqm(dqm, label='Example - Graph Partitioning DQM')

sample = sampleset.first.sample
energy = sampleset.first.energy

color_list = [(random(), random(), random()) for i in range(n_clusters)]
color_map = [color_list[sample[i]] for i in G.nodes]
nx.draw(G, node_color=color_map)
plt.savefig('graph_partition_result.png')
plt.close()

reason: import numpy,import networkx


text:
import math

from keras.layers import Input, LSTM, Dense
from keras.models import Model
from keras import backend as K
from keras.callbacks import EarlyStopping, Callback
from keras.optimizers import Adam, SGD
import matplotlib.pyplot as plt
import sys

import numpy as np
sys.path.append("..")
import config

Z_DIM = config.latent_vector_dim
ACTION_DIM = 2

HIDDEN_UNITS = 256

BATCH_SIZE = 32
EPOCHS = 5

# draw the training loss
class LossHistory(Callback):
    def __init__(self):
        self.losses = {'batch': [], 'epoch': []}
        self.accuracy = {'batch': [], 'epoch': []}
        self.val_loss = {'batch': [], 'epoch': []}
        self.val_acc = {'batch': [], 'epoch': []}

    def on_train_begin(self, logs={}):
        pass

    def on_batch_end(self, batch, logs={}):
        self.losses['batch'].append(logs.get('loss'))
        self.accuracy['batch'].append(logs.get('acc'))
        self.val_loss['batch'].append(logs.get('val_loss'))
        self.val_acc['batch'].append(logs.get('val_acc'))

    def on_epoch_end(self, batch, logs={}):
        self.losses['epoch'].append(logs.get('loss'))
        self.accuracy['epoch'].append(logs.get('acc'))
        self.val_loss['epoch'].append(logs.get('val_loss'))
        self.val_acc['epoch'].append(logs.get('val_acc'))

    def loss_plot(self, loss_type):
        iters = range(len(self.losses[loss_type]))
        plt.figure()
        # acc
        # plt.plot(iters, self.accuracy[loss_type], 'r', label='train acc')
        # loss
        plt.plot(iters, self.losses[loss_type], 'g', label='train loss')
        if loss_type == 'epoch':
            # val_acc
            # plt.plot(iters, self.val_acc[loss_type], 'b', label='val acc')
            # val_loss
            plt.plot(iters, self.val_loss[loss_type], 'k', label='val loss')
        plt.grid(True)
        plt.xlabel(loss_type)
        plt.ylabel('loss')
        plt.legend(loc="upper right")
        plt.show()

class RNN():
    def __init__(self):
        self.models = self._build()
        self.model = self.models[0]
        self.forward = self.models[1]
        self.history = LossHistory()
        self.z_dim = Z_DIM
        self.action_dim = ACTION_DIM
        self.hidden_units = HIDDEN_UNITS

    def _build(self):
        # the model that will be trained
        rnn_x = Input(shape=(None, Z_DIM + ACTION_DIM))
        lstm = LSTM(HIDDEN_UNITS, return_sequences=True, return_state=True)

        lstm_output, _, _ = lstm(rnn_x)
        mdn = Dense(Z_DIM)(lstm_output)

        rnn = Model(rnn_x, mdn)

        # the model used during prediction
        state_input_h = Input(shape=(HIDDEN_UNITS,))
        state_input_c = Input(shape=(HIDDEN_UNITS,))
        state_inputs = [state_input_h, state_input_c]

        _, state_h, state_c = lstm(rnn_x, initial_state=state_inputs)
        forward = Model([rnn_x] + state_inputs, [state_h, state_c])

        optimizer = Adam(lr=0.0001)
        # optimizer = SGD(lr=0.0001, decay=1e-4, momentum=0.9, nesterov=True)
        rnn.compile(loss='mean_squared_error', optimizer=optimizer)

        return [rnn, forward]

    def set_weights(self, filepath):
        self.model.load_weights(filepath)

    def train(self, rnn_input, rnn_output,):
        callbacks_list = [self.history]

        self.model.fit(rnn_input, rnn_output,
                       shuffle=True,
                       epochs=EPOCHS,
                       batch_size=BATCH_SIZE,
                       callbacks=callbacks_list)

        self.model.save_weights('./models/rnn_weights_RNN.h5')

    def plot_loss(self):
        self.history.loss_plot('batch')

    def save_weights(self, filepath):
        self.model.save_weights(filepath)

    # predict the output
    def get_output(self, input_data):
        output = self.model.predict(input_data)
        return output

max_issues_repo_name: goldenminerlmg/navbot

max_forks_repo_name: goldenminerlmg/navbot


text:
"""The types used in the package."""
# --------------------------- Import libraries and functions --------------------------
from typing import List, Tuple, TypeVar, Union

from multiprocessing import pool

import numpy as np
import pandas as pd
from numpy import typing as npt

# ---------------------------------- types definition ---------------------------------
ListFloat = List[float]
ListInt = List[int]
TupleFloat = Tuple[float]
TupleInt = Tuple[int]
ArrayFloat = npt.NDArray[np.float_]
ArrayInt = npt.NDArray[np.int_]
PoolType = TypeVar("PoolType", bound=pool.Pool)
NumArrayLike = TypeVar(
    "NumArrayLike",
    bound=Union[
        pd.Series, ArrayInt, ArrayFloat, ListFloat, ListInt, TupleFloat, TupleInt
    ],
)

max_stars_repo_name: fossabot/pycvcqv

max_issues_repo_name: fossabot/pycvcqv

max_forks_repo_name: fossabot/pycvcqv

reason: import numpy,from numpy


text:
import numpy    as np
import unittest as ut

from qfactor.gates import RxGate

class TestRxGateRepr ( ut.TestCase ):

    def test_rxgate_repr_1 ( self ):
        gate = RxGate( 0., 0 )
        self.assertEqual( repr( gate ), "(0,): Rx(0.0)" )

    def test_rxgate_repr_2 ( self ):
        gate = RxGate( 2., 1 )
        self.assertEqual( repr( gate ), "(1,): Rx(2.0)" )

if __name__ == '__main__':
    ut.main()

max_stars_repo_name: edyounis/QFactor

max_issues_repo_name: edyounis/qFactor

max_forks_repo_name: edyounis/qFactor