text:
#!/usr/bin/env python

import os
import sys
import tarfile
import shutil

import numpy as np
import scipy.linalg
import logging

from contextlib import contextmanager

from pymatgen.core.periodic_table import DummySpecie
from pymatgen.io.gaussian import GaussianInput, GaussianOutput

from tinydb import TinyDB, Query

@contextmanager
def cd(run_path):
    """
    Temporarily work in another directory, creating it if necessary.
    """
    home = os.getcwd()
    if not os.path.isdir(run_path):
        os.makedirs(run_path)
    os.chdir(run_path)
    yield
    os.chdir(home)

def add_dummy_atoms(mol, ring_atom_ids):
    data = np.array([mol[i].coords for i in ring_atom_ids])
    centre = np.average(data, axis=0)

    # fit the atomic positions to a plane using least squares method, taken
    # from: https://gist.github.com/amroamroamro/1db8d69b4b65e8bc66a6#file-curve_fitting-py-L27
    A = np.c_[data[:, 0], data[:, 1], np.ones(data.shape[0])]
    C, _, _, _ = scipy.linalg.lstsq(A, data[:, 2])

    # calculate the normal unit vector
    normal = np.array(np.cross([1, 0, C[0]], [0, 1, C[1]]))
    unit = normal/scipy.linalg.norm(normal)

    # add the dummy atoms above and below the ring
    mol.append(DummySpecie('X-Bq'), centre + unit)
    mol.append(DummySpecie('X-Bq'), centre - unit)
    return mol

# define calculation settings
functional = 'b3lyp'
dieze_tag = '#p'
basis_set = '6-311++G(d,p)'
link0 = {'%oldchk': 'chkpt.chk', '%mem': '58GB', '%nprocshared': '24'}
td_params = {'integral': '(acc2e=12)', 'td': '(50-50)', 'guess': 'read'}
tda_params = {'integral': '(acc2e=12)', 'tda': '(50-50)', 'guess': 'read'}
nmr_params = {'integral': '(acc2e=12)', 'guess': 'read', 'nmr': ''}

# define the rings used to calculate NICS
six_mem_a = [0, 1, 2, 3, 4, 5]
six_mem_b = [1, 2, 6, 7, 8, 9]
five_mem_a = [0, 1, 9, 13, 12]
five_mem_b = [2, 3, 10, 11, 6]

index = int(sys.argv[1])

# use absolute path so we don't loose track of the db when changing directory
db = TinyDB(os.path.abspath('structures.json'))
query = Query()
compound = db.get(query.index == index)
rin = GaussianInput.from_dict(compound['input'])
directory = rin.title

def calculate_properties(rin, directory):

    with cd(directory):

        # found that the relax often crashed or didn't finish, so will restart
        # the calculation in these cases
        try:
            rout = GaussianOutput('relax.log')
            if not rout.properly_terminated:
                logging.info('restarting {}'.format(directory))
                route = rout.route
                route['integral'] = '(acc2e=12)'
                rout.to_input('relax.com', cart_coords=True,
                              route_parameters=route)
                os.system('g09 < relax.com > relax.log')
        except IOError:
            # relaxation hasn't been run yet
            logging.info('started processing {}'.format(directory))
            rin.write_file('relax.com', cart_coords=True)
            os.system('g09 < relax.com > relax.log')
        except IndexError, AttributeError:
            # the relax calculation used the wrong header, fix this and restart
            rin = GaussianInput.from_file('relax.com')
            rin.route_parameters['integral'] = '(acc2e=12)'
            rin.write_file('relax.com', cart_coords=True)
            os.system('g09 < relax.com > relax.log')

        rout = GaussianOutput('relax.log')
        if not rout.properly_terminated:
            logging.error('{} relaxation did not terminate correctly'.format(
                          directory))
            return 0

        # do the TD-DFT calculation
        tdin = GaussianInput(rout.final_structure, charge=0, title=rin.title,
                             functional=functional, spin_multiplicity=1,
                             basis_set=basis_set, dieze_tag=dieze_tag,
                             link0_parameters=link0, route_parameters=td_params)
        td_done = False
        if os.path.isfile('td.log'):
            tdout = GaussianOutput('td.log')
            if tdout.properly_terminated:
                td_done = True

        if not td_done:
            tdin.write_file('td.com', cart_coords=True)
            os.system('g09 < td.com > td.log')
            tdout = GaussianOutput('td.log')
            if not tdout.properly_terminated:
                logging.error('{} TD-DFT did not terminate correctly'.format(
                              directory))
                return 0

        # do the TD-DFT calculation w. Tamm-Dancoff approx
        tdain = GaussianInput(rout.final_structure, charge=0, title=rin.title,
                              spin_multiplicity=1, functional=functional,
                              basis_set=basis_set, dieze_tag=dieze_tag,
                              link0_parameters=link0,
                              route_parameters=tda_params)
        tda_done = False
        if os.path.isfile('tda.log'):
            tdaout = GaussianOutput('tda.log')
            if tdaout.properly_terminated:
                tda_done = True

        if not tda_done:
            tdain.write_file('tda.com', cart_coords=True)
            os.system('g09 < tda.com > tda.log')
            tdaout = GaussianOutput('tda.log')
            if not tdaout.properly_terminated:
                logging.error('{} TDA-DFT did not terminate correctly'.format(
                              directory))
                return 0

        # add the dummy atoms for the NICS(1)_zz calculations
        mol_nics = rout.final_structure
        mol_nics = add_dummy_atoms(mol_nics, six_mem_a)
        mol_nics = add_dummy_atoms(mol_nics, six_mem_b)
        mol_nics = add_dummy_atoms(mol_nics, five_mem_a)
        mol_nics = add_dummy_atoms(mol_nics, five_mem_b)

        # run NICS on the ground state and triplet state
        nicssin = GaussianInput(mol_nics, charge=0, title=rin.title,
                                spin_multiplicity=1, functional=functional,
                                basis_set=basis_set, dieze_tag=dieze_tag,
                                link0_parameters=link0,
                                route_parameters=nmr_params)
        nicssin.write_file('nics_singlet.com', cart_coords=True)
        # work around as pymatgen does not allow Bq as an element
        os.system("sed -i 's/X-Bq0+/Bq/g' nics_singlet.com")
        os.system('g09 < nics_singlet.com > nics_singlet.log')
        # can't have NICS job completion check due to above mentioned pmg bug

        nicstin = GaussianInput(mol_nics, charge=0, title=rin.title,
                                spin_multiplicity=3, functional=functional,
                                basis_set=basis_set, dieze_tag=dieze_tag,
                                link0_parameters=link0,
                                route_parameters=nmr_params)
        nicstin.write_file('nics_triplet.com', cart_coords=True)
        os.system("sed -i 's/X-Bq0+/Bq/g' nics_triplet.com")
        os.system('g09 < nics_triplet.com > nics_triplet.log')
    logging.info('finished processing {}'.format(directory))
    return 1

with cd('calculations'):
    logging.basicConfig(filename='calculations.log', level=logging.DEBUG)

    tar_file = '{}.tar.gz'.format(directory)
    error_tar_file = '{}_error.tar.gz'.format(directory)

    if os.path.isfile(tar_file):
        sys.exit()
    elif os.path.isfile(error_tar_file) and os.path.isdir(directory):
        os.remove(error_tar_file)
    elif os.path.isfile(error_tar_file):
        with tarfile.open(error_tar_file) as tar:
            tar.extractall()
        os.remove(error_tar_file)

    finished = calculate_properties(rin, directory)
    filename = tar_file if finished == 1 else error_tar_file

    os.remove(os.path.join(directory, 'chkpt.chk'))
    with tarfile.open(filename, "w:gz") as tar:
        tar.add(directory)
    shutil.rmtree(directory)

max_stars_repo_name: SMTG-UCL/singlet-fission-screening

max_issues_repo_name: SMTG-UCL/singlet-fission-screening

max_forks_repo_name: SMTG-UCL/singlet-fission-screening


text:
import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt

def add_layer(inputs, in_size, out_size, n_layer, activation_function=None):
    layer_name = "layer" + str(n_layer)
    with tf.name_scope("layers"):
        with tf.name_scope("weights"):
            weights = tf.Variable(tf.random_normal([in_size, out_size]), name="W")
            tf.histogram_summary(layer_name + "/weights", weights)
        with tf.name_scope("biases"):
            biases = tf.Variable(tf.zeros([1, out_size]) + 0.1, "b")
            tf.histogram_summary(layer_name + "/biases", biases)
        with tf.name_scope("Wx_plus_b"):
            wx_plus_b = tf.matmul(inputs, weights) + biases

        if activation_function:
            outputs = activation_function(wx_plus_b)
        else:
            outputs = wx_plus_b
        tf.histogram_summary(layer_name + "/outputs", outputs)
        return outputs

# Make up some real data.
x_data = np.linspace(-1, 1, 300)[:, np.newaxis]
noise = np.random.normal(0, 0.05, size=x_data.shape)
y_data = np.square(x_data) - 0.5 + noise

# define place hold for inputs to network.
with tf.name_scope("input"):
    xs = tf.placeholder(tf.float32, [None, 1], name="x_input")
    ys = tf.placeholder(tf.float32, [None, 1], name="y_input")

l1 = add_layer(xs, 1, 10, n_layer=0, activation_function=tf.nn.relu)
# Linear regression, no activation or linear activation.
prediction = add_layer(l1, 10, 1, n_layer=1, activation_function=None)

with tf.name_scope("loss"):
    loss = tf.reduce_mean(tf.reduce_sum(tf.square(ys - prediction), reduction_indices=[1]))
    tf.scalar_summary("loss", loss)

with tf.name_scope("train"):
    optimizer = tf.train.AdamOptimizer(learning_rate=0.01)
    train_step = optimizer.minimize(loss)

#
init = tf.initialize_all_variables()

sess = tf.Session()

merged = tf.merge_all_summaries()
writer = tf.train.SummaryWriter("./graphs/", sess.graph)

sess.run(init)
#
# plt.ion()
# figure, axes = plt.subplots(2, 1, sharex=False, sharey=False)
# top = axes[0]
# bot = axes[1]
# top.scatter(x_data, y_data, c='b')
#
# perc = 1
#
#
# def vis():
#     loss_val = sess.run(loss, feed_dict={xs: x_data, ys: y_data})
#     prediction_val = sess.run(prediction, feed_dict={xs: x_data})
#     lines = top.plot(x_data, prediction_val, c='r')
#     bot.scatter(i, loss_val, c='k')
#     plt.pause(0.05)
#     top.lines.remove(lines[0])
#
for i in range(1001):
    sess.run(train_step, feed_dict={xs: x_data, ys: y_data})
    if i % 50 == 0:
        result = sess.run(merged, feed_dict={xs: x_data, ys: y_data})
        writer.add_summary(result, i)
else:
    # vis()
    pass
#
#
# plt.ioff()
# plt.show()

ext: py

max_stars_repo_name: nitred/tf-examples

max_issues_repo_name: nitred/tf-examples

max_forks_repo_name: nitred/tf-examples

reason: import numpy


text:
# Copyright 2018 Mamy André-Ratsimbazafy. All rights reserved.

import pandas as pd
import numpy as np
from src.instrumentation import logspeed

@logspeed
def fte_cyclic_time(train, test, y, db_conn, folds, cache_file):
  def _trans(df, table):
    query = f"""
    select
      --case(WEEKDAY_APPR_PROCESS_START)
      --  when "SATURDAY" THEN 1
      --  when "SUNDAY" THEN 1
      --  else 0
      --end isOccidentalWeekend,
      case(WEEKDAY_APPR_PROCESS_START)
        when "MONDAY" THEN 0
        when "TUESDAY" THEN 1
        when "WEDNESDAY" THEN 2
        when "THURSDAY" THEN 3
        when "FRIDAY" THEN 4
        when "SATURDAY" THEN 5
        when "SUNDAY" THEN 6
        else "NA"
      end dayOfWeek,
      HOUR_APPR_PROCESS_START
    from
      {table}
    order by
      SK_ID_CURR ASC;
    """

    df[[
      # 'isOccidentalWeekend',
      'dayOfWeek',
      'HOUR_APPR_PROCESS_START']] = pd.read_sql_query(query, db_conn)
    # df['cos_dayOfWeek'] = np.cos(df['dayOfWeek'] / 7 * 2 * np.pi)
    # df['sin_dayOfWeek'] = np.sin(df['dayOfWeek'] / 7 * 2 * np.pi)
    # df['cos_hour_start'] = np.cos(df['HOUR_APPR_PROCESS_START'] / 24 * 2 * np.pi)
    # df['sin_hour_start'] = np.sin(df['HOUR_APPR_PROCESS_START'] / 24 * 2 * np.pi)

  _trans(train, "application_train")
  _trans(test, "application_test")

  return train, test, y, db_conn, folds, cache_file

max_stars_repo_name: mratsim/home-credit-default-risk

max_issues_repo_name: mratsim/home-credit-default-risk

max_forks_repo_name: mratsim/home-credit-default-risk


text:
import numpy as np
from pandas import read_csv
from glue.config import data_factory
from glue.core import Data

import pandas as pd
import gzip

from pyxlma.lmalib.io.read import lmafile

# class lmafile:
#     def __init__(self,filename):
#         self.file = filename
#
#     def datastart(self):
#         with gzip.open(self.file,'rt') as f:
#             for line_no, line in enumerate(f):
#                     if line.rstrip() == "*** data ***":
#                         break
#             f.close()
#             return line_no
#
#     def readfile(self):
#         lmad = pd.read_csv(self.file,delim_whitespace=True,header=None,
#                                           compression='gzip',skiprows=self.datastart()+1)
#         columns = ['time','lat','lon','alt','chi','p','mask']
#         lmad.columns = columns
#         return lmad

def is_lma_dataset(filename, **kwargs):
    return '.dat' in filename

@data_factory('LMA data file', priority=10000,
              identifier=is_lma_dataset)
def read_lma_data(filename):
    lma = lmafile(filename)
    df = lma.readfile()
    data = Data()
    for column in df.columns:
        data[column] = df[column]
    # data['distance'] = np.hypot(data['x'], data['y'])
    return data

# def is_glm_dataset(filename, **kwargs):
#     return 'OR_GLM' in filename
#
# @data_factory('GLM gridded image', priority=10000,
#               identifier=is_glm_dataset)
# def read_glm_image(filename):
#     # Load the GLM data
#     glm = open_glm_time_series([filename])
#     x_1d = glm.x
#     y_1d = glm.y
#
#     # Convert the 1D fixed grid coordinates to 2D lon, lat
#
#     from lmatools.grid.fixed import get_GOESR_coordsys
#     x,y = np.meshgrid(x_1d, y_1d) # Two 2D arrays of fixed grid coordinates
#     nadir = -75.0 # Really should come from file.
#     geofixCS, grs80lla = get_GOESR_coordsys(nadir)
#     z=np.zeros_like(x)
#     lon,lat,alt=grs80lla.fromECEF(*geofixCS.toECEF(x,y,z))
#     lon.shape = x.shape
#     lat.shape = y.shape
#
#     # Add the 2D arrays back to the original dataset and save to disk.
#     # This doesn't have the CF standard name metadata or unit information,
#     # but that isn't too hard to add.
#
#     import xarray as xr
#     glm['longitude'] = xr.DataArray(lon, dims=('y', 'x'))
#     glm['latitude'] = xr.DataArray(lat, dims=('y', 'x'))

max_stars_repo_name: vbalderdash/xlma-python

max_issues_repo_name: vbalderdash/xlma-python

max_forks_repo_name: vbalderdash/xlma-python


text:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from __future__ import division
from __future__ import print_function

import numpy as np
import pandas as pd
import copy
import random
from .pytorch_utils import count_parameters
from ...utils import (
    get_or_create_path,
)
from ...log import get_module_logger

import torch
import torch.nn as nn
import torch.optim as optim

from ...model.base import Model
from ...data.dataset import DatasetH
from ...data.dataset.handler import DataHandlerLP

class IPPE(Model):
    """IPPE Model

    Parameters
    ----------
    d_feat : int
        input dimension for each time step
    metric: str
        the evaluate metric used in early stop
    optimizer : str
        optimizer name
    GPU : str
        the GPU ID(s) used for training
    """

    def __init__(
        self,
        d_feat=6,
        hidden_size=64,
        num_layers=2,
        dropout=0.0,
        n_epochs=200,
        batch_size=100,
        early_stop=5,
        loss="mse",
        output_dim=5,
        lr=5e-4,
        use_pretrain_pair=True,
        use_pretrain_point=True,
        pretrain_epoch=20,
        pretrain_lr=2e-4,
        point_enhance=True,
        pair_enhance=True,
        point_dim=10,
        pair_dim=10,
        iter_num=1,
        evaluate_time=3,
        inference_time=10,
        metric="loss",
        optimizer="adam",
        GPU=0,
        seed=None,
        **kwargs
    ):
        # Set logger.
        self.logger = get_module_logger("IPPE")
        self.logger.info("IPPE pytorch version...")

        # set hyper-parameters.
        self.d_feat = d_feat
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.n_epochs = n_epochs
        self.batch_size = batch_size
        self.early_stop = early_stop
        self.loss = loss
        self.output_dim = output_dim
        self.lr = lr
        self.use_pretrain_pair = use_pretrain_pair
        self.use_pretrain_point = use_pretrain_point
        self.pretrain_epoch = pretrain_epoch
        self.pretrain_lr = pretrain_lr
        self.point_enhance = point_enhance
        self.pair_enhance = pair_enhance
        self.point_dim = point_dim
        self.pair_dim = pair_dim
        self.iter_num = iter_num
        self.evaluate_time = evaluate_time
        self.inference_time = inference_time
        self.optimizer = optimizer
        self.metric = metric
        self.device = torch.device("cuda:%d" % (GPU) if torch.cuda.is_available() else "cpu")
        self.seed = seed

        self.logger.info(
            "IPPE parameters setting:"
            "\nd_feat : {}"
            "\nhidden_size : {}"
            "\nnum_layers : {}"
            "\ndropout : {}"
            "\nn_epochs : {}"
            "\nbatch_size : {}"
            "\nearly_stop : {}"
            "\nloss_type : {}"
            "\noutput_dim : {}"
            "\nlr : {}"
            "\nuse_pretrain_pair : {}"
            "\nuse_pretrain_point : {}"
            "\npretrain_epoch : {}"
            "\npretrain_lr : {}"
            "\npoint_enhance : {}"
            "\npair_enhance : {}"
            "\npoint_dim : {}"
            "\npair_dim : {}"
            "\niter_num : {}"
            "\ninference_time : {}"
            "\noptimizer : {}"
            "\nseed : {}".format(
                d_feat,
                hidden_size,
                num_layers,
                dropout,
                n_epochs,
                batch_size,
                early_stop,
                loss,
                output_dim,
                lr,
                use_pretrain_pair,
                use_pretrain_point,
                pretrain_epoch,
                pretrain_lr,
                point_enhance,
                pair_enhance,
                point_dim,
                pair_dim,
                iter_num,
                inference_time,
                optimizer,
                seed,
            )
        )

        if self.seed is not None:
            np.random.seed(self.seed)
            torch.manual_seed(self.seed)
        self.point_model = GRU_POINT(
            d_feat=self.d_feat,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
        ).to(self.device)
        self.pair_model = GRU_PAIR(
            d_feat=self.d_feat*2,
            hidden_size=self.hidden_size,
            num_layers=self.num_layers,
            dropout=self.dropout,
        ).to(self.device)
        self.model = MUTUAL_ENHANCE(device=self.device,
                                    point_input=self.hidden_size, \
                                    pair_input=self.hidden_size, \
                                    point_dim=self.point_dim, \
                                    pair_dim=self.pair_dim, \
                                    point_enhance=self.point_enhance, \
                                    pair_enhance=self.pair_enhance).to(self.device)
        for param in self.model.parameters():
            param.requires_grad = True

        param_list = list(self.model.parameters()) + list(self.point_model.parameters()) + list(self.pair_model.parameters())

        self.logger.info("IPPE model size: {:.4f} MB".format(count_parameters(self.model)))

        if optimizer.lower() == "adam":
            self.train_optimizer = optim.Adam(param_list, lr=self.lr)
            if self.use_pretrain_point:
                self.point_pretrain_optimizer = optim.Adam(self.point_model.parameters(), lr=self.pretrain_lr)
            if self.use_pretrain_pair:
                self.pair_pretrain_optimizer = optim.Adam(self.pair_model.parameters(), lr=self.pretrain_lr)
        elif optimizer.lower() == "gd":
            self.train_optimizer = optim.SGD(param_list, lr=self.lr)
            if self.use_pretrain_point:
                self.point_pretrain_optimizer = optim.SGD(self.point_model.parameters(), lr=self.pretrain_lr)
            if self.use_pretrain_pair:
                self.pair_pretrain_optimizer = optim.SGD(self.pair_model.parameters(), lr=self.pretrain_lr)
        else:
            raise NotImplementedError("optimizer {} is not supported!".format(optimizer))

        self.fitted = False

    @property
    def use_gpu(self):
        return self.device != torch.device("cpu")

    def mse(self, pred, label):
        loss = (pred - label) ** 2
        return torch.mean(loss)

    def loss_fn(self, pred, label):
        mask = ~torch.isnan(label)

        if self.loss == "mse":
            return self.mse(pred[mask], label[mask])

        raise ValueError("unknown loss `%s`" % self.loss)

    def metric_fn(self, pred, label):

        mask = torch.isfinite(label)

        if self.metric == "" or self.metric == "loss":
            return -self.loss_fn(pred[mask], label[mask])

        raise ValueError("unknown metric `%s`" % self.metric)

    def pair_data_process(self, feature, label):

        feature = np.array(feature)
        label = np.array(label)
        stock_num = feature.shape[0]

        pair_label = []
        if label.shape[0]:
            vec_1 = np.repeat(label, stock_num).reshape(stock_num, stock_num)
            vec_2 = np.transpose(vec_1)
            pair_label = (vec_1 - vec_2).reshape(-1)

        loc_1 = np.repeat(range(stock_num), stock_num).reshape(stock_num, stock_num)
        loc_2 = np.transpose(loc_1)
        loc_1 = loc_1.reshape([-1])
        loc_2 = loc_2.reshape([-1])

        return feature[loc_1], feature[loc_2], pair_label

    def train_epoch(self, x_train, y_train):

        self.model.train(True)
        self.point_model.train(True)
        self.pair_model.train(True)

        days = x_train.index.get_level_values(0).unique()
        total_day = days.shape[0]
        day_index_seq = np.array(range(total_day))
        random.shuffle(day_index_seq)

        for _, day_index in enumerate(day_index_seq):
            # prepare feature and label
            day = days[day_index]
            feature = np.array(x_train.loc[day])
            label = np.array(y_train.loc[day]).reshape([-1])
            stock_num = feature.shape[0]

            loc = np.random.choice(stock_num, self.batch_size)
            feature = feature[loc]
            label = label[loc]

            pair_feature1, pair_feature2, pair_label = self.pair_data_process(feature, label)

            point_feature = torch.tensor(feature, dtype=torch.float, device=self.device)
            point_label = torch.tensor(label, dtype=torch.float, device=self.device)
            pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
            pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)
            pair_label = torch.tensor(pair_label, dtype=torch.float, device=self.device)

            # transfer feature to embedding
            point = self.point_model(point_feature)
            pair = self.pair_model(pair_feature1, pair_feature2)
            point = self.model.transform_point(point)
            pair = self.model.transform_pair(pair)

            point_loss = 0
            pair_loss = 0
            if self.point_enhance:
                point_pred = self.model.point2pred(point)
                point_loss += self.loss_fn(point_pred.view(-1), point_label.view(-1))
            if self.pair_enhance:
                pair_pred = self.model.pair2pred(pair).view(-1)
                pair_loss += self.loss_fn(pair_pred.view(-1), pair_label.view(-1))

            for _ in range(self.iter_num):
                point, pair = self.model(point, pair)

                if self.point_enhance:
                    point_pred = self.model.point2pred(point)
                    point_loss += self.loss_fn(point_pred.view(-1), point_label.view(-1))
                if self.pair_enhance:
                    pair_pred = self.model.pair2pred(pair).view(-1)
                    pair_loss += self.loss_fn(pair_pred.view(-1), pair_label.view(-1))

            loss = (point_loss + pair_loss)/((1 + self.iter_num)*self.iter_num/2)

            self.train_optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_value_(self.point_model.parameters(), 3.0)
            self.train_optimizer.step()

    def test_epoch(self, data_x, data_y):
        self.model.eval()
        self.point_model.eval()
        self.pair_model.eval()

        scores = []
        losses = []

        days = data_x.index.get_level_values(0).unique()
        total_day = days.shape[0]
        day_index_seq = np.array(range(total_day))

        for _ in range(self.evaluate_time):
            for _, day_index in enumerate(day_index_seq):
                # prepare feature and label
                day = days[day_index]
                feature = np.array(data_x.loc[day])
                label = np.array(data_y.loc[day]).reshape([-1])
                stock_num = feature.shape[0]

                loc = np.random.choice(stock_num, self.batch_size)
                feature = feature[loc]
                label = label[loc]

                pair_feature1, pair_feature2, pair_label = self.pair_data_process(feature, label)
                point_feature = torch.tensor(feature, dtype=torch.float, device=self.device)
                point_label = torch.tensor(label, dtype=torch.float, device=self.device)
                pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
                pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)
                pair_label = torch.tensor(pair_label, dtype=torch.float, device=self.device)

                with torch.no_grad():
                    # transfer feature to embedding
                    point = self.point_model(point_feature)
                    pair = self.pair_model(pair_feature1, pair_feature2)
                    point = self.model.transform_point(point)
                    pair = self.model.transform_pair(pair)

                    for _ in range(self.iter_num):
                        point, pair = self.model(point, pair)
                    point_pred = self.model.point2pred(point)

                    loss = self.loss_fn(point_pred, point_label)
                    losses.append(loss.item())
                    score = self.metric_fn(point_pred, point_label)
                    scores.append(score.item())

        return np.mean(losses), np.mean(scores)

    def pretrain_point_model(
        self,
        x_train,
        y_train,
        batch_size=2000,
    ):

        x_train_values = x_train.values
        y_train_values = np.squeeze(y_train.values)

        self.point_model.train()

        for step in range(self.pretrain_epoch):
            self.logger.info("Epoch%d:", step)
            indices = np.arange(len(x_train_values))
            np.random.shuffle(indices)

            for i in range(len(indices))[:: batch_size]:

                if len(indices) - i < batch_size:
                    break

                feature = torch.from_numpy(x_train_values[indices[i : i + batch_size]]).float().to(self.device)
                label = torch.from_numpy(y_train_values[indices[i : i + batch_size]]).float().to(self.device)

                point_embedding = self.point_model(feature)
                pred = self.point_model.fc_out(point_embedding).squeeze()
                loss = self.loss_fn(pred, label)

                self.point_pretrain_optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_value_(self.point_model.parameters(), 3.0)
                self.point_pretrain_optimizer.step()

    def pretrain_pair_model(
        self,
        x_train,
        y_train,
    ):

        days = x_train.index.get_level_values(0).unique()
        total_day = days.shape[0]
        day_index_seq = np.array(range(total_day))
        random.shuffle(day_index_seq)

        for step in range(self.pretrain_epoch):
            self.logger.info("Epoch%d:", step)
            for _, day_index in enumerate(day_index_seq):
                # prepare feature and label
                day = days[day_index]
                feature = np.array(x_train.loc[day])
                label = np.array(y_train.loc[day]).reshape([-1])
                stock_num = feature.shape[0]

                loc = np.random.choice(stock_num, self.batch_size)
                feature = feature[loc]
                label = label[loc]

                pair_feature1, pair_feature2, pair_label = self.pair_data_process(feature, label)

                pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
                pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)
                pair_label = torch.tensor(pair_label, dtype=torch.float, device=self.device)

                pair_embedding = self.pair_model(pair_feature1, pair_feature2)
                pair_pred = self.pair_model.fc_out(pair_embedding).squeeze()
                loss = self.loss_fn(pair_pred, pair_label)

                self.pair_pretrain_optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_value_(self.pair_model.parameters(), 3.0)
                self.pair_pretrain_optimizer.step()

    def fit(
        self,
        dataset: DatasetH,
        evals_result=dict(),
        save_path=None,
    ):

        df_train, df_valid, _ = dataset.prepare(
            ["train", "valid", "test"],
            col_set=["feature", "label"],
            data_key=DataHandlerLP.DK_L,
        )
        save_path = get_or_create_path(save_path)
        stop_steps = 0
        best_score = -np.inf
        best_epoch = 0
        evals_result["valid"] = []

        self.logger.info("training...")
        self.fitted = True

        if self.use_pretrain_point:
            self.logger.info("pretrain point model...")
            self.pretrain_point_model(df_train["feature"], df_train["label"])

        if self.use_pretrain_pair:
            self.logger.info("pretrain pair model...")
            self.pretrain_pair_model(df_train["feature"], df_train["label"])

        for step in range(self.n_epochs):
            self.logger.info("Epoch%d:", step)
            self.train_epoch(df_train["feature"], df_train["label"])
            _, val_score = self.test_epoch(df_valid["feature"], df_valid["label"])
            self.logger.info("valid %.6f" % (val_score))
            evals_result["valid"].append(val_score)

            if val_score > best_score:
                best_score = val_score
                stop_steps = 0
                best_epoch = step
                best_param = [copy.deepcopy(self.model.state_dict()),\
                                copy.deepcopy(self.point_model.state_dict()),\
                                copy.deepcopy(self.pair_model.state_dict())]

            else:
                stop_steps += 1
                if stop_steps >= self.early_stop:
                    self.logger.info("early stop")
                    break

        self.logger.info("best score: %.6lf @ %d" % (best_score, best_epoch))
        self.model.load_state_dict(best_param[0])
        self.point_model.load_state_dict(best_param[1])
        self.pair_model.load_state_dict(best_param[2])
        torch.save(best_param[0], save_path+'ippe')
        torch.save(best_param[1], save_path+'point_model')
        torch.save(best_param[2], save_path+'pair_model')

    def predict(self, dataset: DatasetH):
        if not self.fitted:
            raise ValueError("model is not fitted yet!")

        test_data = dataset.prepare("test", col_set=["feature", "label"], data_key=DataHandlerLP.DK_I)
        x_test = test_data["feature"]
        y_test = test_data["label"]
        test_data.to_pickle("~/tmp_test_data.pkl")
        self.model.eval()
        self.point_model.eval()
        self.pair_model.eval()

        preds = []
        indexes = []
        days = x_test.index.get_level_values(0).unique()
        total_day = days.shape[0]

        for day_index in range(total_day):
            # prepare feature and label
            day = days[day_index]
            daily_feature = np.array(x_test.loc[day])
            daily_index = x_test.loc(axis=0)[day,:].index
            stock_num = daily_feature.shape[0]

            for _ in range(self.inference_time):
                loc = np.random.choice(stock_num, self.batch_size)
                feature = daily_feature[loc]
                currend_index = daily_index[loc]

                pair_feature1, pair_feature2, _ = self.pair_data_process(feature, [])
                point_feature = torch.tensor(feature, dtype=torch.float, device=self.device)
                pair_feature1 = torch.tensor(pair_feature1, dtype=torch.float, device=self.device)
                pair_feature2 = torch.tensor(pair_feature2, dtype=torch.float, device=self.device)

                with torch.no_grad():
                    # transfer feature to embedding
                    point = self.point_model(point_feature)
                    pair = self.pair_model(pair_feature1, pair_feature2)
                    point = self.model.transform_point(point)
                    pair = self.model.transform_pair(pair)

                    for _ in range(self.iter_num):
                        point, pair = self.model(point, pair)
                    point_pred = self.model.point2pred(point)

                preds.append(point_pred.detach().cpu().numpy().reshape([-1]))
                indexes.append(currend_index)

        res_df = pd.Series(np.concatenate(preds), index=np.concatenate(indexes))
        res_df.index = pd.MultiIndex.from_tuples(res_df.index)
        res_df.index.names = ['datetime', 'instrument']
        res_df = res_df.groupby(level=[0,1]).apply(lambda x:x.mean())
        preds = pd.Series(np.zeros(x_test.shape[0]), index=test_data.index)
        preds.loc[res_df.index] = np.array(res_df)

        return preds

class GRU_POINT(nn.Module):

    def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0):
        super().__init__()

        self.rnn = nn.GRU(
            input_size=d_feat,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
        )
        self.fc_out = nn.Linear(hidden_size, 1)

        self.d_feat = d_feat

    def forward(self, x):
        # x: [N, F*T]
        x = x.reshape(len(x), self.d_feat, -1) # [N, F, T]
        x = x.permute(0, 2, 1) # [N, T, F]
        out, _ = self.rnn(x)
        return out[:, -1, :].squeeze()

class GRU_PAIR(nn.Module):

    def __init__(self, d_feat=6, hidden_size=64, num_layers=2, dropout=0.0):
        super().__init__()

        self.rnn = nn.GRU(
            input_size=d_feat*2,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
        )
        self.fc_out = nn.Linear(hidden_size, 1)
        self.d_feat = d_feat

    def forward(self, x, y):
        # x: [N, F*T]
        x = x.reshape(len(x), self.d_feat, -1) # [N, F, T]
        x = x.permute(0, 2, 1) # [N, T, F]
        y = y.reshape(len(y), self.d_feat, -1) # [N, F, T]
        y = y.permute(0, 2, 1) # [N, T, F]

        input = torch.cat((x, y), 2)
        out, _ = self.rnn(input)
        return out[:, -1, :].squeeze()

class MUTUAL_ENHANCE(nn.Module):
    def __init__(self, device, point_input, pair_input, point_dim=10, pair_dim=10, point_enhance=True, pair_enhance=True):

        super().__init__()
        self.device = device
        self.point_dim = point_dim
        self.pair_dim = pair_dim
        self.point_enhance = point_enhance
        self.pair_enhance = pair_enhance
        self.transform_point = nn.Linear(point_input, point_dim)
        self.transform_pair = nn.Linear(pair_input, pair_dim)
        self.point2pred = nn.Linear(point_dim, 1)
        self.pair2pred = nn.Linear(pair_dim, 1)

        # prepare lamda block
        self.transferion_point = torch.nn.Parameter(torch.rand(self.point_dim, self.point_dim*self.pair_dim).cuda(), \
                                             requires_grad=True)

        self.transferion_pair = torch.nn.Parameter(torch.rand(self.point_dim, self.pair_dim*self.point_dim).cuda(), \
                                             requires_grad=True)

        self.lamda_point_func = nn.Sequential(
            nn.Linear(3*self.point_dim+3, 16),
            nn.ReLU(),
            nn.Linear(16,16),
            nn.ReLU(),
            nn.Linear(16,1),
            nn.Sigmoid())

        self.omega_point_func = nn.Sequential(
            nn.Linear(3*self.point_dim + 2*self.pair_dim + 2, 32),
            nn.ReLU(),
            nn.Linear(32,32),
            nn.ReLU(),
            nn.Linear(32,16),
            nn.ReLU(),
            nn.Linear(16,1))

        self.lamda_pair_func = nn.Sequential(
            nn.Linear(2*self.point_dim + 3*self.pair_dim + 2,16),
            nn.ReLU(),
            nn.Linear(16,16),
            nn.ReLU(),
            nn.Linear(16,1),
            nn.Sigmoid())

        self.softmax = torch.nn.Softmax(dim = 1)

    def forward(self, point, pair):
        # data prepare

        stock_num = point.shape[0]
        point_1 = point.expand(stock_num, stock_num, point.shape[1])
        point_2 = torch.transpose(point_1,0,1)
        pair = pair.view(stock_num, stock_num, -1)

        pair_1 = pair
        pair_2 = torch.transpose(pair_1,0,1)

        if self.point_enhance:
            # calculate point estimation embedding
            correction_in_1 = pair_1.view(stock_num, stock_num, self.pair_dim, 1)
            correction_in_2 = point_1.view(stock_num, stock_num, 1, self.point_dim)
            tmp = torch.matmul(correction_in_2, self.transferion_point)
            tmp = tmp.view(stock_num, stock_num, self.point_dim, self.pair_dim)
            self.point_estimation = torch.matmul(tmp, correction_in_1).view(stock_num, stock_num, self.point_dim)
            self.point_estimation_pred = self.point2pred(self.point_estimation)

            # omega in
            self.point_estimation_mean = torch.sum(self.point_estimation, dim=1)/stock_num
            self.point_estimation_var = torch.std(self.point_estimation, dim=1)
            self.point_estimation_pred_mean = torch.sum(self.point_estimation_pred, dim=1)/stock_num
            self.point_estimation_pred_var = torch.std(self.point_estimation_pred, dim=1)

            self.omega_in = torch.cat((pair_1, # embedding(i,j)
                                        pair_2, # embedding(j,i)
                                        point_1, # embedding(j)
                                        point_2, # embedding(i)
                                        self.point_estimation,
                                        self.point_estimation_pred_mean.expand(stock_num, stock_num).view(stock_num, stock_num, -1),
                                        self.point_estimation_pred_var.expand(stock_num, stock_num).view(stock_num, stock_num, -1),
                                  ), 2)

            # omega out
            self.omega_point = self.omega_point_func(self.omega_in)
            self.omega_point = self.softmax(self.omega_point)

            # calculate aggerated esetimation embedding
            self.agg_point_estimation = torch.sum(self.omega_point*self.point_estimation, dim=1)
            # self.agg_point_estimation = torch.sum(self.point_estimation, dim=1)/stock_num # average

            # lambda in
            agg_point_estimation_pred = self.point_estimation_pred*self.omega_point
            agg_point_estimation_pred_mean = torch.sum(agg_point_estimation_pred, dim=1)/stock_num
            agg_point_estimation_pred_var = torch.std(agg_point_estimation_pred, dim=1)
            omega_std = torch.std(self.omega_point, dim=1)
            distance = torch.abs(self.agg_point_estimation - point)
            self.lamda_in_point = torch.cat((point,
                              self.agg_point_estimation.view(stock_num, -1),
                              agg_point_estimation_pred_mean.view(stock_num, -1),
                              agg_point_estimation_pred_var.view(stock_num, -1),
                              omega_std.view(stock_num, -1),
                              distance.view(stock_num, -1)
                             ), 1)
            # lambda out
            self.lamda_point = self.lamda_point_func(self.lamda_in_point)
            # update
            point = (1-self.lamda_point)*point + self.lamda_point*self.agg_point_estimation

        if self.pair_enhance:
            # calculate pair estimation embedding
            correction_in_1 = point_1.view(stock_num, stock_num, 1, self.point_dim)
            correction_in_2 = point_2.view(stock_num, stock_num, self.point_dim, 1)
            tmp = torch.matmul(correction_in_1, self.transferion_pair)
            tmp = tmp.view(stock_num, stock_num, self.pair_dim, self.point_dim)
            self.pair_estimation = torch.matmul(tmp, correction_in_2).view(stock_num, stock_num, self.pair_dim)
            self.pair_estimation_pred = self.point2pred(self.pair_estimation)

            # lambda in
            self.pair_estimation_mean = torch.sum(self.pair_estimation, dim=1)/stock_num
            self.pair_estimation_var = torch.std(self.pair_estimation, dim=1)
            self.pair_estimation_pred_mean = torch.sum(self.pair_estimation_pred, dim=1)/stock_num
            self.pair_estimation_pred_var = torch.std(self.pair_estimation_pred, dim=1)

            self.lamda_in_pair = torch.cat((pair_1, # embedding(i,j)
                                        pair_2, # embedding(j,i)
                                        point_1, # embedding(j)
                                        point_2, # embedding(i)
                                        self.pair_estimation,
                                        self.pair_estimation_pred_mean.expand([stock_num, stock_num]).view(stock_num, stock_num, -1),
                                        self.pair_estimation_pred_var.expand([stock_num, stock_num]).view(stock_num, stock_num, -1),
                                  ), 2)

            # lambda out
            self.lambda_pair = self.lamda_pair_func(self.lamda_in_pair)
            # update
            pair = (1-self.lambda_pair)*pair + self.lambda_pair*self.pair_estimation

        return point, pair


text:
from __future__ import division, print_function
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import *
import numpy as np
from keras.datasets import mnist
from keras.utils.np_utils import to_categorical
from keras.optimizers import SGD
from sklearn.linear_model import LogisticRegression

def get_cnn(n_classes):
    model = Sequential()
    model.add(Convolution2D(32, 3, 3, activation='relu', input_shape=(28,28,1)))
    model.add(Convolution2D(32, 3, 3, activation='relu'))
    model.add(MaxPooling2D())
    model.add(Convolution2D(64, 3, 3, activation='relu'))
    model.add(Convolution2D(64, 3, 3, activation='relu'))
    model.add(MaxPooling2D())
    model.add(Flatten())
    model.add(Dropout(0.2))
    model.add(Dense(128, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(n_classes, activation='softmax'))
    model.compile(optimizer='Adam',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def get_mlp(n_classes):
    model = Sequential()
    model.add(Dense(128, activation='relu',input_shape=(784,)))
    model.add(Dropout(0.5))
    model.add(Dense(128, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(n_classes, activation='softmax'))
    model.compile(optimizer='Adam',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def get_mlp_leek(n_classes):
    model = Sequential()
    model.add(Dense(160, activation='tanh',input_shape=(784,)))
    model.add(Dense(160, activation='tanh'))
    model.add(Dense(160, activation='tanh'))
    model.add(Dense(160, activation='tanh'))
    model.add(Dense(160, activation='tanh'))
    model.add(Dense(n_classes, activation='softmax'))
    ## Attempts to match h2o documentation as closely as possible
    ## https://cran.r-project.org/web/packages/h2o/h2o.pdf
    opt = SGD(lr=0.005, momentum=0.0,decay=0.99, nesterov=True)
    model.compile(optimizer=opt,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def get_leekassso_predictors(x,y,n_preds):
    eps = 1e-6
    ## This really makes me appreciate dplyr!
    x_0 = x[np.where(y==0)]
    x_1 = x[np.where(y==1)]
    x0_bar = x_0.mean(axis=0)
    x1_bar = x_1.mean(axis=0)
    x0_sigma = x_0.std(axis=0)
    x1_sigma = x_1.std(axis=0)
    sigma_pool = np.sqrt( (x0_sigma**2)/len(x_0) + (x1_sigma**2)/len(x_1) ) + eps
    all_t = np.abs((x0_bar - x1_bar)/sigma_pool)
    top_10 = np.argsort(-all_t)[0:n_preds]
    return top_10

## Load the data ##
(X_train, y_train), (X_test, y_test) = mnist.load_data()

## Set the number of training samples to use ##
train_sizes = [10,20,40,60,80]
batch_size = 1

X_train = np.expand_dims(X_train, axis=3)
X_test = np.expand_dims (X_test, axis=3)

# Get training data for 0 and 1
inds = np.where((y_train == 0) | (y_train == 1))[0]

X_train_small = X_train[inds]
X_train_small = X_train_small[:train_size]
y_train_small = y_train[inds]
y_train_small = y_train_small[:train_size]

# Get test data for 0 and 1 only
inds_test = np.where((y_test == 0) | (y_test == 1))
X_test_01 = X_test[inds_test]
y_test_01 = y_test[inds_test]

# split into validation and test sets
test_start_ind = int(np.floor((len(X_test_01)/2)))
X_test_small = X_test_01[:test_start_ind]
X_final_test = X_test_01[test_start_ind:]

y_test_small = y_test_01[:test_start_ind]
y_final_test = y_test_01[test_start_ind:]

# one hot labels
one_hot_train = to_categorical(y_train_small, 2)
one_hot_test = to_categorical(y_test_small, 2)
one_hot_final_test = to_categorical(y_final_test, 2)

n_classes = one_hot_train.shape[1]

folds = 5
evals = np.zeros((len(train_sizes*folds),3))
index = 0
for train_size in train_sizes:
    for i in range(5):
        fold_inds = np.random.choice(inds,train_size)
        X_train_fold = X_train[fold_inds]
        y_train_fold = y_train[fold_inds]
        one_hot_fold = to_categorical(y_train_fold, 2)
        model = get_cnn(n_classes)
        model.fit(X_train_fold, one_hot_fold, nb_epoch=200)
        score = (model.evaluate(X_final_test, one_hot_final_test, batch_size=128))[1]
        evals[index,0] = train_size
        evals[index,1] = i
        evals[index,2] = score
        index += 1

np.savetxt('cnn.csv',evals)

X_test_flat = X_test_small.reshape((len(X_test_small),784))
X_final_test_flat = X_final_test.reshape((len(X_final_test),784))

folds = 5
evals = np.zeros((len(train_sizes*folds),3))
index = 0
for train_size in train_sizes:
    for i in range(5):
        fold_inds = np.random.choice(inds,train_size)
        X_train_fold = X_train[fold_inds].reshape(train_size,784)
        y_train_fold = y_train[fold_inds]
        one_hot_fold = to_categorical(y_train_fold, 2)
        model = get_mlp(n_classes)
        model.fit(X_train_fold, one_hot_fold, nb_epoch=200,
                        validation_data=[X_test_flat, one_hot_test])
        score = (model.evaluate(X_final_test_flat, one_hot_final_test, batch_size=batch_size))[1]
        evals[index,0] = train_size
        evals[index,1] = i
        evals[index,2] = score
        index += 1

np.savetxt('mlp.csv',evals)

folds = 5
evals = np.zeros((len(train_sizes*folds),3))
index = 0
for train_size in train_sizes:
    for i in range(5):
        fold_inds = np.random.choice(inds,train_size)
        X_train_fold = X_train[fold_inds].reshape(train_size,784)
        y_train_fold = y_train[fold_inds]
        one_hot_fold = to_categorical(y_train_fold, 2)
        model = get_mlp_leek(n_classes)
        model.fit(X_train_fold, one_hot_fold, nb_epoch=20,batch_size=1)
        score = (model.evaluate(X_final_test_flat, one_hot_final_test, batch_size=256))[1]
        evals[index,0] = train_size
        evals[index,1] = i
        evals[index,2] = score
        index += 1

np.savetxt('mlp_leek.csv',evals)

folds = 5
evals = np.zeros((len(train_sizes*folds),3))
index = 0
for train_size in train_sizes:
    for i in range(5):
        fold_inds = np.random.choice(inds,train_size)
        X_train_fold = X_train[fold_inds].reshape(train_size,784)
        y_train_fold = y_train[fold_inds]
        leekasso_preds = get_leekassso_predictors(X_train_fold,y_train_fold,10)
        X_leek = X_train_fold[:,leekasso_preds]
        model = LogisticRegression(C=1e6) ## no L2 penalty
        model.fit(X_leek, y_train_fold)
        score = model.score(X_final_test_flat[:,leekasso_preds],y_final_test)
        evals[index,0] = train_size
        evals[index,1] = i
        evals[index,2] = score
        index += 1

np.savetxt('leekasso.csv',evals)


text:
import numpy as np
import os
from pymol import cmd

##### FOR analysis of DRRAFTER simulated benchmark ######
####### (VERY specific to this use case) #######
## run in pymol, because I was doing some additional visual analysis...

systems = ['1b7f',
	'1dfu',
	'1p6v',
        '1jbs',
        '1wpu',
        '1wsu',
        '2asb',
        '2bh2',
        '2qux',
        '3bx2']

native_pdbs = { '1b7f': '1b7f_bound.pdb',
		'1jbs': '1jbs_bound_orig.pdb',
		'1wpu': '1wpu_bound_NEW.pdb',
		'1wsu': '1wsu_bound.pdb',
		'2asb': '2asb_bound.pdb',
		'2bh2': '2bh2_bound_full.pdb',
		'2qux': '2qux_bound.pdb',
		'3bx2': '3bx2_bound.pdb',
		'1p6v': '1p6v_bound.pdb',
		'1dfu': '1dfu_bound.pdb' }

resolutions = ['3A', '5A', '7A']

start_dir = os.getcwd()

for sys in systems:
	base_name=sys
	for reso in resolutions:
		print sys, reso, '\n\n'
		os.chdir( start_dir )
		# change to the directory
		run_dir = sys
		if sys == '1jbs':
			run_dir += '/add_BP/'
		run_dir += '/%s/' %(reso)

		os.chdir( run_dir )

		rmsd_sele='resn A+U+G+C'

		for s1 in range(1,11):
			cmd.load('%s_fold_and_dock.out.%d.pdb' %(base_name,s1), 's%d' %(s1))
		# and load the native
		cmd.load('%s/native_pdbs/%s' %(start_dir, native_pdbs[sys]),'native')

		## get the RMSD to native
		#rmsds_to_native = []
		#for s1 in range(1,11):
		#	cmd.align('not resn A+U+G+C and s%d' %(s1), 'native')
		#	rms = cmd.rms_cur( '%s and s%d' %(rmsd_sele, s1), 'native')
		#	rmsds_to_native.append( rms )

		f = open('pairwise_rmsds_top_10.txt','w')
		pairwise_rmsds = []
		for s1 in range(1,11):
			for s2 in range(s1+1,11):
				# align based on the protein
				cmd.align( 'not resn A+U+G+C and s%d' %(s1), 'not resn A+U+G+C and s%d' %(s2) )
				# get only the RNA rmsd (without alignment)
				rms = cmd.rms_cur( '%s and s%d' %(rmsd_sele, s1), '%s and s%d' %(rmsd_sele, s2) )
				pairwise_rmsds.append( rms )
				print s1, s2, rms
				f.write('%d %d %0.3f\n' %(s1,s2,rms))

		pairwise_rmsds = np.array( pairwise_rmsds )
		avg_rmsd = np.mean( pairwise_rmsds )
		#rmsds_to_native = np.array( rmsds_to_native )
		#min_rmsd_to_native = np.min( rmsds_to_native )
		print "Mean rmsd ", avg_rmsd
		#print "Min rmsd to native ", min_rmsd_to_native
		f.write('Mean rmsd %0.3f\n' %(avg_rmsd))
		#f.write('Min rmsd to native %0.3f\n' %(min_rmsd_to_native))

		f.close()

		# close models
		cmd.delete('native')
		for s1 in range(1,11):
			cmd.delete('s%d' %(s1))


text:
from copy import copy

import numpy as np
import progressbar
from bpdb import set_trace

from pntddf.env import setup_env
from pntddf.results import (comm_savings, plot_b, plot_b_dot, plot_residuals,
                            plot_rover_state_errors, plot_test, plot_time,
                            plot_trajectory)

config_file = "../config/sim.config"
# config_file = "../config/1d.config"
# config_file = "../config/demo.config"

env = setup_env(config_file)

print("Running Simulation")
increments = 1000

with progressbar.ProgressBar(max_value=env.MAX_TIME) as bar:
    for t in np.linspace(env.MAX_TIME / increments, env.MAX_TIME, increments):
        bar.update(np.round(t))
        env.run(until=t)

print("Generating Plots")
plots = [
    plot_b,
    plot_b_dot,
    plot_rover_state_errors,
    plot_trajectory,
    plot_residuals,
    plot_time,
]

agents_to_plot = copy(env.agents)
if env.centralized:
    agents_to_plot += [env.agent_centralized]

counter = 0
with progressbar.ProgressBar(max_value=len(agents_to_plot) * len(plots)) as bar:
    for agent in agents_to_plot:
        for plot in plots:
            df_state = agent.estimator.get_state_log_df()
            df_meas = agent.estimator.get_residuals_log_df()
            plot(env, agent, df_state, df_meas)
            counter += 1
            bar.update(counter)

# comm_savings(env, env.agents)
plot_test(env, env.agents[0])

max_stars_repo_name: COHRINT/pnt-ddf

max_issues_repo_name: COHRINT/pnt-ddf

max_forks_repo_name: COHRINT/pnt-ddf


text:
# Classify the MNIST digits using a one nearest neighbour classifier and Euclidean distance

import numpy as np

DEFAULT_TEST_SIZE = None
DEFAULT_BUCKET_SIZE = 1000
MINIMUM_BUCKET_SIZE = 100

def PredictandError(trainims, trainlabs, testims, testlabs):

    from sklearn.metrics.pairwise import euclidean_distances
    Distances = euclidean_distances(trainims, testims, squared=True)

    predictions = trainlabs[np.argmin(Distances, axis=0)]
    return np.not_equal(predictions, testlabs).sum()

def run_main(TestSize, BucketSize):

    from keras.datasets import mnist
    (TrainIms, TrainLabels), (TestIms, TestLabels) = mnist.load_data()

    TrainIms = TrainIms.reshape(TrainIms.shape[0], -1)
    TestIms = TestIms.reshape(TestIms.shape[0], -1)

    if TestSize is DEFAULT_TEST_SIZE:
        TestSize = len(TrainLabels)
    BucketSize = min(BucketSize, TestSize)

    import time
    t0 = time.time()

    # error: 3.0900000000000025
    # Time taken: 19.416749477386475
    errors = [
        PredictandError(TrainIms, TrainLabels,
                        TestIms[i:min(i+BucketSize, TestSize)],
                        TestLabels[i:min(i+BucketSize, TestSize)])
        for i in range(0, TestSize, BucketSize)
    ]

    # error: 3.090000000000004
    # Time taken: 80.45218276977539
    #from sklearn.neighbors import KNeighborsClassifier
    #clf = KNeighborsClassifier(n_neighbors=1, algorithm='brute', p=2, n_jobs=-1)
    #testims = TrainIms[:TestSize]
    #testlabs = TrainLabels[:TestSize]
    #clf.fit(testims, testlabs)
    #errors = [(1 - clf.score(testims, testlabs)) * 100]

    t1 = time.time()

    print("error: {:.3f} %".format(sum(errors) / TestSize * 100))
    print("Time taken: {:.3f} sec".format(t1 - t0))

def parse_args():

    import argparse
    parser = argparse.ArgumentParser()

    parser.add_argument('--test-size', type=int, default=DEFAULT_TEST_SIZE,
                        help='size of the test images')
    parser.add_argument('--bucket-size', type=int, default=DEFAULT_BUCKET_SIZE,
                        help='size of the test bucket')
    args = parser.parse_args()

    assert (args.test_size is DEFAULT_TEST_SIZE or args.test_size >= DEFAULT_BUCKET_SIZE), \
        'Too small test size (< {}).'.format(DEFAULT_BUCKET_SIZE)
    assert (args.bucket_size >= MINIMUM_BUCKET_SIZE), \
        'Too small bucket size (< {}).'.format(MINIMUM_BUCKET_SIZE)

    return args.test_size, args.bucket_size

if __name__ == '__main__':
    test_size, bucket_size = parse_args()
    run_main(test_size, bucket_size)

max_stars_repo_name: chofchof/probml

max_issues_repo_name: chofchof/probml

max_forks_repo_name: chofchof/probml


text:
#!/usr/bin/env python
# coding: utf-8

import scipy.optimize
import json
import numpy as np
import re
import sys
import math
import argparse
from collections import defaultdict
from pprint import pprint

def float_list(s):
    return [float(x) for x in s.split(",")] if s else []

parser = argparse.ArgumentParser()
parser.add_argument("--transport-power-cost", type=float, default=50.0,
    help="added power cost for transport per conveyor/pipeline of mined resource")
parser.add_argument("--drone-battery-cost", type=float, default=0.5,
    help="added battery cost for drone transport per conveyor/pipeline of mined resource")
parser.add_argument("--machine-penalty", type=float, default=2000.0,
    help="objective penalty per machine built")
parser.add_argument("--conveyor-penalty", type=float, default=0.0,
    help="objective penalty per conveyor belt needed")
parser.add_argument("--pipeline-penalty", type=float, default=0.0,
    help="objective penalty per pipeline needed")
parser.add_argument("--power-shard-penalty-ratio", type=float, default=0.6,
    help="objective penalty per power shard used, specified as ratio of machine penalty")
parser.add_argument("--extra-miner-clocks", type=float_list, default=[],
    help="extra clock choices for miners, specified as decimals")
parser.add_argument("--extra-manufacturer-clocks", type=float_list, default=[0.25, 0.5, 0.75],
    help="extra clock choices for manufacturers, specified as decimals")
parser.add_argument("--allow-waste", action="store_true",
    help="allow accumulation of nuclear waste and other unsinkable items")
parser.add_argument("--show-unused", action="store_true",
    help="show unused LP columns (coeff 0) in the optimization result")
parser.add_argument("--xlsx-report", type=str, default="Report.xlsx",
    help="path to xlsx report output")
parser.add_argument("--xlsx-sheet-suffix", type=str, default="",
    help="suffix to add to xlsx sheet names")
args = parser.parse_args()

### Constants ###

# Common
STACK_SIZES = {
    "SS_HUGE": 500,
    "SS_BIG": 200,
    "SS_MEDIUM": 100,
    "SS_SMALL": 50,
    "SS_ONE": 1,
    "SS_FLUID": 50000,
}
MACHINE_POWER_SHARD_LIMIT = 3
EPSILON = 1e-9

# Logistics
CONVEYOR_BELT_CLASS = "Build_ConveyorBeltMk5_C"
PIPELINE_CLASS = "Build_PipelineMK2_C"

# Resource extraction
MINER_CLASS = "Build_MinerMk3_C"
OIL_EXTRACTOR_CLASS = "Build_OilPump_C"
WATER_EXTRACTOR_CLASS = "Build_WaterPump_C"
RESOURCE_WELL_EXTRACTOR_CLASS = "Build_FrackingExtractor_C"
RESOURCE_WELL_PRESSURIZER_CLASS = "Build_FrackingSmasher_C"

# Sink
SINK_CLASS = "Build_ResourceSink_C"

# Water
WATER_CLASS = "Desc_Water_C"

# Nuclear power
NUCLEAR_WASTE_MAPPINGS = {
    "Desc_NuclearFuelRod_C": "Desc_NuclearWaste_C",
    "Desc_PlutoniumFuelRod_C": "Desc_PlutoniumWaste_C",
}

# Geothermal power
GEOTHERMAL_GENERATOR_CLASS = "Build_GeneratorGeoThermal_C"
GEYSER_CLASS = "Desc_Geyser_C"

# Resource map
PURITY_MULTIPLIERS = {
    "impure": 0.5,
    "normal": 1.0,
    "pure": 2.0,
}
POWER_SLUG_SHARDS = {
    "greenSlugs": 1,
    "yellowSlugs": 2,
    "purpleSlugs": 5,
}
RESOURCE_MAPPINGS = {
    "Desc_LiquidOilWell_C": "Desc_LiquidOil_C",
    "Desc_SAM_C": None,  # exclude
}

# Miscellaneous
BIOMASS_GENERATOR_CLASS = "Build_GeneratorBiomass_C"
BATTERY_CLASS = "Desc_Battery_C"
ADDITIONAL_ITEMS = {
    "Desc_PlutoniumWaste_C": {
        "class": "Desc_PlutoniumWaste_C",
        "display_name": "Plutonium Waste",
        "form": "RF_SOLID",
        "points": 0,
        "stack_size": STACK_SIZES["SS_HUGE"],
        "energy": 0.0,
    },
}
ADDITIONAL_DISPLAY_NAMES = {
    GEYSER_CLASS: "Geyser",
}

docs_path = r"Docs.json"
map_info_path = r"MapInfo.json"

with open(docs_path, "r", encoding="utf-16") as f:
    docs_raw = json.load(f)

class_entries = {}
class_types = {}

for fg_entry in docs_raw:
    class_type = re.sub(r"Class'/Script/FactoryGame.(\w+)'", r"\1", fg_entry["NativeClass"])
    class_type_list = []
    for class_entry in fg_entry["Classes"]:
        class_name = class_entry["ClassName"]
        if class_name in class_entries:
            print(f"WARNING: ignoring duplicate class {class_name}")
        else:
            class_entries[class_name] = class_entry
            class_type_list.append(class_entry)
    class_types[class_type] = class_type_list

### Parsing helpers ###

def parse_paren_list(s):
    if not s:
        return None
    assert(s.startswith("(") and s.endswith(")"))
    s = s[1:-1]
    if not s:
        return []
    else:
        return s.split(",")

def find_class_name(s):
    m = re.search(r"\.\w+", s)
    if m is None:
        raise ValueError(f"could not find class name in: {s}")
    return m[0][1:]

def parse_class_list(s):
    l = parse_paren_list(s)
    if l is None:
        return l
    return [find_class_name(x) for x in l]

def find_item_amounts(s):
    for m in re.finditer(r"\(ItemClass=([^,]+),Amount=(\d+)\)", s):
        yield (find_class_name(m[1]), int(m[2]))

### Misc constants ###

CONVEYOR_BELT_LIMIT = 0.5 * float(class_entries[CONVEYOR_BELT_CLASS]["mSpeed"])
PIPELINE_LIMIT = 60000.0 * float(class_entries[PIPELINE_CLASS]["mFlowLimit"])
SINK_POWER_CONSUMPTION = float(class_entries[SINK_CLASS]["mPowerConsumption"])

print(f"CONVEYOR_BELT_LIMIT: {CONVEYOR_BELT_LIMIT}")
print(f"PIPELINE_LIMIT: {PIPELINE_LIMIT}")
print(f"SINK_POWER_CONSUMPTION: {SINK_POWER_CONSUMPTION}")

### Miners ###

def parse_miner(entry):
    if entry["ClassName"] == RESOURCE_WELL_PRESSURIZER_CLASS:
        extractor = class_entries[RESOURCE_WELL_EXTRACTOR_CLASS]
    else:
        extractor = entry

    return {
        "class": entry["ClassName"],
        "display_name": entry["mDisplayName"],
        "power_consumption": float(entry["mPowerConsumption"]),
        "power_consumption_exponent": float(entry["mPowerConsumptionExponent"]),
        "min_clock": float(entry["mMinPotential"]),
        "max_clock_base": float(entry["mMaxPotential"]),
        "max_clock_per_power_shard": float(entry["mMaxPotentialIncreasePerCrystal"]),
        "rate": 60.0 / float(extractor["mExtractCycleTime"]) * float(extractor["mItemsPerCycle"]),
        "only_allow_certain_resources": (extractor["mOnlyAllowCertainResources"] == "True"),
        "allowed_resource_forms": parse_paren_list(extractor["mAllowedResourceForms"]),
        "allowed_resources": parse_class_list(extractor["mAllowedResources"]),
    }

miners = {}
for name in (MINER_CLASS, OIL_EXTRACTOR_CLASS, WATER_EXTRACTOR_CLASS, RESOURCE_WELL_PRESSURIZER_CLASS):
    miners[name] = parse_miner(class_entries[name])

# pprint(miners)

### Manufacturers ###

def parse_manufacturer(entry):
    return {
        "class": entry["ClassName"],
        "display_name": entry["mDisplayName"],
        "power_consumption": float(entry["mPowerConsumption"]),
        "power_consumption_exponent": float(entry["mPowerConsumptionExponent"]),
        "min_clock": float(entry["mMinPotential"]),
        "max_clock_base": float(entry["mMaxPotential"]),
        "max_clock_per_power_shard": float(entry["mMaxPotentialIncreasePerCrystal"]),
    }

manufacturers = {}

for entry in class_types["FGBuildableManufacturer"]:
    manufacturer = parse_manufacturer(entry)
    manufacturer["is_variable_power"] = False
    manufacturers[entry["ClassName"]] = manufacturer

for entry in class_types["FGBuildableManufacturerVariablePower"]:
    manufacturer = parse_manufacturer(entry)
    manufacturer["is_variable_power"] = True
    manufacturers[entry["ClassName"]] = manufacturer

# pprint(manufacturers)

### Recipes ###

def parse_recipe(entry):
    recipe_manufacturer = None
    for manufacturer in parse_class_list(entry["mProducedIn"]) or []:
        if manufacturer in manufacturers:
            recipe_manufacturer = manufacturer
            break

    # we are only considering automatable recipes
    if recipe_manufacturer is None:
        return None

    rate = 60.0 / float(entry["mManufactoringDuration"])
    def item_rates(key):
        return [(item, rate * amount) for (item, amount) in find_item_amounts(entry[key])]

    vpc_constant = float(entry["mVariablePowerConsumptionConstant"])
    vpc_factor = float(entry["mVariablePowerConsumptionFactor"])

    return {
        "class": entry["ClassName"],
        "display_name": entry["mDisplayName"],
        "manufacturer": recipe_manufacturer,
        "inputs": item_rates("mIngredients"),
        "outputs": item_rates("mProduct"),
        "variable_power_consumption": vpc_constant + 0.5 * vpc_factor,
    }

recipes = {}
for entry in class_types["FGRecipe"]:
    recipe = parse_recipe(entry)
    if recipe is not None:
        recipes[entry["ClassName"]] = recipe

# pprint(recipes)

### Items ###

def parse_item(entry):
    points = int(entry["mResourceSinkPoints"])
    return {
        "display_name": entry["mDisplayName"],
        "form": entry["mForm"],
        "points": int(entry["mResourceSinkPoints"]),
        "stack_size": STACK_SIZES[entry["mStackSize"]],
        "energy": float(entry["mEnergyValue"]),
    }

items = {}

# any items not contained in Docs.json
items.update(ADDITIONAL_ITEMS)

for class_type in [
    "FGItemDescriptor",
    "FGItemDescriptorBiomass",
    "FGItemDescriptorNuclearFuel",
    "FGResourceDescriptor",
    "FGEquipmentDescriptor",
    "FGConsumableDescriptor",
]:
    for entry in class_types[class_type]:
        item = parse_item(entry)
        if class_type == "FGItemDescriptorNuclearFuel":
            item["nuclear_waste"] = NUCLEAR_WASTE_MAPPINGS[entry["ClassName"]]
            item["nuclear_waste_amount"] = float(entry["mAmountOfWaste"])
        items[entry["ClassName"]] = item

# pprint(items)

### Generators ###

generators = {}

def parse_generator(entry):
    power_production = float(entry["mPowerProduction"])
    return {
        "display_name": entry["mDisplayName"],
        "fuel_classes": parse_class_list(entry["mDefaultFuelClasses"]),
        "power_production": power_production,
        "power_production_exponent": float(entry["mPowerProductionExponent"]),
        "requires_supplemental": (entry["mRequiresSupplementalResource"] == "True"),
        "supplemental_to_power_ratio": float(entry["mSupplementalToPowerRatio"]),
    }

def parse_geothermal_generator(entry):
    # unclear why mVariablePowerProductionConstant=0 in the json;
    # it's set to 100.0f in the header, which we will hardcode here
    return {
        "display_name": entry["mDisplayName"],
        "power_production": 100.0 + 0.5 * float(entry["mVariablePowerProductionFactor"]),
    }

# coal and fuel generators
for entry in class_types["FGBuildableGeneratorFuel"]:
    # exclude biomass generator
    if entry["ClassName"] == BIOMASS_GENERATOR_CLASS:
        continue

    generators[entry["ClassName"]] = parse_generator(entry)

# nuclear power plant
for entry in class_types["FGBuildableGeneratorNuclear"]:
    generators[entry["ClassName"]] = parse_generator(entry)

# geothermal generator (special case)
geothermal_generator = parse_geothermal_generator(class_entries[GEOTHERMAL_GENERATOR_CLASS])

# pprint(generators)

### Resources ###

with open(map_info_path, "r") as f:
    map_info_raw = json.load(f)

map_info = {}

for tab in map_info_raw["options"]:
    if "tabId" in tab:
        map_info[tab["tabId"]] = tab["options"]

TOTAL_POWER_SHARDS = 0
for slug_type in map_info["power_slugs"][0]["options"]:
    TOTAL_POWER_SHARDS += POWER_SLUG_SHARDS[slug_type["layerId"]] * len(slug_type["markers"])

print(f"TOTAL_POWER_SHARDS: {TOTAL_POWER_SHARDS}")

resources = {}
geysers = {}

def parse_and_add_node_type(node_type):
    if "type" not in node_type:
        return

    item = node_type["type"]
    if item in RESOURCE_MAPPINGS:
        item = RESOURCE_MAPPINGS[item]

    if item is None:
        return

    output = geysers if item == GEYSER_CLASS else resources

    for node_purity in node_type["options"]:
        purity = node_purity["purity"]
        nodes = node_purity["markers"]
        if not nodes:
            continue
        sample_node = nodes[0]
        if "core" in sample_node:
            # resource well satellite nodes, map them to cores
            for node in nodes:
                subtype = find_class_name(node["core"])
                resource_id = f"{item}|{subtype}"
                if resource_id not in output:
                    output[resource_id] = {
                        "resource_id": resource_id,
                        "item": item,
                        "subtype": subtype,
                        "multiplier": 0.0,
                        "count": 1,
                        "is_limited": True,
                        "is_resource_well": True,
                        "num_satellites": 0,
                    }
                output[resource_id]["multiplier"] += PURITY_MULTIPLIERS[purity]
                output[resource_id]["num_satellites"] += 1
        else:
            # normal nodes, add directly
            subtype = purity
            resource_id = f"{item}|{subtype}"
            assert(resource_id not in output)
            output[resource_id] = {
                "resource_id": resource_id,
                "item": item,
                "subtype": subtype,
                "multiplier": PURITY_MULTIPLIERS[purity],
                "count": len(nodes),
                "is_limited": True,
                "is_resource_well": False,
            }

for node_type in map_info["resource_nodes"]:
    parse_and_add_node_type(node_type)

for node_type in map_info["resource_wells"]:
    parse_and_add_node_type(node_type)

resources[WATER_CLASS] = {
    "resource_id": f"{WATER_CLASS}:extractor",
    "item": WATER_CLASS,
    "subtype": "extractor",
    "multiplier": 1,
    "is_limited": False,
    "is_resource_well": False,
}

# pprint(resources)
# pprint(geysers)

### LP setup ###

class LPColumn(dict):
    def __init__(self, *args, display_info=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.display_info = display_info

lp_columns = {}
lp_equalities = {}
lp_lower_bounds = {}

def get_power_consumption(machine, clock=1.0, recipe=None):
    power_consumption = machine["power_consumption"]
    if recipe is not None and machine.get("is_variable_power", False):
        power_consumption += recipe["variable_power_consumption"]

    return power_consumption * (clock ** machine["power_consumption_exponent"])

def get_miner_for_resource(resource):
    item_class = resource["item"]
    item = items[item_class]
    candidates = []
    for miner_class, miner in miners.items():
        if ((resource["is_resource_well"]) == (miner_class == RESOURCE_WELL_PRESSURIZER_CLASS)
            and item["form"] in miner["allowed_resource_forms"]
            and (not miner["only_allow_certain_resources"] or item_class in miner["allowed_resources"])):
            candidates.append(miner_class)
    if not candidates:
        raise RuntimeError(f"could not find miner for resource {item_class}")
    elif len(candidates) > 1:
        raise RuntimeError(f"more than one miner for resource {item_class}: {candidates}")
    return candidates[0]

def get_form_conveyance_limit(form):
    if form == "RF_SOLID":
        return CONVEYOR_BELT_LIMIT
    elif form == "RF_LIQUID" or form == "RF_GAS":
        return PIPELINE_LIMIT
    else:
        assert(False)

def get_max_overclock(machine):
    return machine["max_clock_base"] + MACHINE_POWER_SHARD_LIMIT * machine["max_clock_per_power_shard"]

def get_conveyance_limit_clock(item, rate):
    conveyance_limit = get_form_conveyance_limit(item["form"])
    return math.floor(1000000 * conveyance_limit / rate) / 1000000

def get_max_miner_clock(miner, resource, rate):
    max_overclock = get_max_overclock(miner)

    if resource["is_resource_well"]:
        return max_overclock

    item_class = resource["item"]
    item = items[item_class]
    return min(max_overclock, get_conveyance_limit_clock(item, rate))

def get_max_manufacturer_clock(manufacturer, recipe):
    max_clock = get_max_overclock(manufacturer)

    for (item_class, rate) in recipe["inputs"]:
        max_clock = min(max_clock, get_conveyance_limit_clock(items[item_class], rate))

    for (item_class, rate) in recipe["outputs"]:
        max_clock = min(max_clock, get_conveyance_limit_clock(items[item_class], rate))

    return max_clock

def get_power_shards_needed(machine, clock):
    return max(0, math.ceil((clock - machine["max_clock_base"]) / machine["max_clock_per_power_shard"]))

def get_item_display_name(item_class):
    if item_class in items:
        return items[item_class]["display_name"]
    else:
        return ADDITIONAL_DISPLAY_NAMES[item_class]

def add_lp_column(column, type_, name, display_name=None, machine_name=None, subtype=None, clock=None):
    tokens = [type_, name]
    if subtype is not None:
        tokens.append(subtype)
    if clock is not None:
        clock_percent = 100.0 * clock
        tokens.append(f"{clock_percent}")
    column_id = "|".join(tokens)
    display_info = {
        "type": type_,
        "display_name": display_name or name,
        "machine_name": machine_name,
        "subtype": subtype,
        "clock": clock,
    }
    lp_columns[column_id] = LPColumn(column, display_info=display_info)

for resource_id, resource in resources.items():
    item_class = resource["item"]
    item = items[item_class]

    miner_class = get_miner_for_resource(resource)
    miner = miners[miner_class]

    rate = miner["rate"] * resource["multiplier"]
    min_clock = miner["min_clock"]
    max_clock_base = miner["max_clock_base"]
    max_clock = get_max_miner_clock(miner, resource, rate)

    resource_var = f"resource|{resource_id}"
    item_var = f"item|{item_class}"

    clock_choices = {max_clock_base, max_clock}
    for clock in args.extra_miner_clocks:
        clock = min(max_clock, max(min_clock, clock))
        clock_choices.add(clock)

    for clock in sorted(clock_choices):
        column = {
            item_var: clock * rate,
            "power_consumption": get_power_consumption(miner, clock=clock),
            "machines": 1 + (resource["num_satellites"] if resource["is_resource_well"] else 0),
        }

        if resource["is_limited"]:
            column[resource_var] = -1

        power_shards = get_power_shards_needed(miner, clock)
        if power_shards > 0:
            column["power_shard_usage"] = power_shards

        add_lp_column(
            column,
            type_="miner",
            name=resource_id,
            display_name=item["display_name"],
            machine_name=miner["display_name"],
            subtype=resource["subtype"],
            clock=clock,
        )

    if resource["is_limited"]:
        lp_lower_bounds[resource_var] = -resource["count"]

    lp_equalities[item_var] = 0.0

for recipe_class, recipe in recipes.items():
    manufacturer_class = recipe["manufacturer"]
    manufacturer = manufacturers[manufacturer_class]

    min_clock = manufacturer["min_clock"]
    max_clock_base = manufacturer["max_clock_base"]
    max_clock = get_max_manufacturer_clock(manufacturer, recipe)

    # let's not allow manufacturer OC by default, but it can be specified via option
    clock_choices = {min_clock, max_clock_base}
    for clock in args.extra_manufacturer_clocks:
        clock = min(max_clock, max(min_clock, clock))
        clock_choices.add(clock)

    for clock in sorted(clock_choices):
        column = {
            "power_consumption": get_power_consumption(manufacturer, clock=clock, recipe=recipe),
            "machines": 1,
        }

        for (item_class, rate) in recipe["inputs"]:
            item_var = f"item|{item_class}"
            column[item_var] = column.get(item_var, 0.0) - clock * rate
            lp_equalities[item_var] = 0.0

        for (item_class, rate) in recipe["outputs"]:
            item_var = f"item|{item_class}"
            column[item_var] = column.get(item_var, 0.0) + clock * rate
            lp_equalities[item_var] = 0.0

        power_shards = get_power_shards_needed(manufacturer, clock)
        if power_shards > 0:
            column["power_shard_usage"] = power_shards

        add_lp_column(
            column,
            type_="manufacturer",
            name=recipe_class,
            display_name=recipe["display_name"],
            machine_name=manufacturer["display_name"],
            clock=clock,
        )

for item_class, item in items.items():
    points = item["points"]
    item_var = f"item|{item_class}"

    if not (item["form"] == "RF_SOLID" and points > 0):
        if args.allow_waste:
            add_lp_column(
                {item_var: -1},
                type_="waste",
                name=item_class,
                display_name=item["display_name"],
            )
        continue

    column = {
        item_var: -1,
        "points": points,
        "power_consumption": SINK_POWER_CONSUMPTION / CONVEYOR_BELT_LIMIT,
        "machines": 1 / CONVEYOR_BELT_LIMIT,
    }

    add_lp_column(
        column,
        type_="sink",
        name=item_class,
        display_name=item["display_name"],
    )

    lp_equalities[item_var] = 0.0

for generator_class, generator in generators.items():
    power_production = generator["power_production"]
    for fuel_class in generator["fuel_classes"]:
        fuel = items[fuel_class]
        fuel_rate = 60.0 * power_production / fuel["energy"]
        fuel_var = f"item|{fuel_class}"

        column = {
            fuel_var: -fuel_rate,
            "power_production": power_production,
            "machines": 1,
        }

        if generator["requires_supplemental"]:
            supplemental_class = WATER_CLASS
            supplemental_var = f"item|{supplemental_class}"
            supplemental_rate = 60.0 * power_production * generator["supplemental_to_power_ratio"]
            column[supplemental_var] = -supplemental_rate
            lp_equalities[supplemental_var] = 0.0

        if fuel_class in NUCLEAR_WASTE_MAPPINGS:
            waste_class = NUCLEAR_WASTE_MAPPINGS[fuel_class]
            waste_var = f"item|{waste_class}"
            column[waste_var] = fuel_rate * fuel["nuclear_waste_amount"]
            lp_equalities[waste_var] = 0.0

        add_lp_column(
            column,
            type_="generator",
            name=fuel_class,
            display_name=fuel["display_name"],
            machine_name=generator["display_name"],
            clock=1,
        )

for resource_id, resource in geysers.items():
    resource_var = f"resource|{resource_id}"

    column = {
        resource_var: -1,
        "power_production": geothermal_generator["power_production"] * resource["multiplier"],
        "machines": 1,
    }

    add_lp_column(
        column,
        type_="generator",
        name=resource_id,
        display_name=get_item_display_name(GEYSER_CLASS),
        machine_name=geothermal_generator["display_name"],
        subtype=resource["subtype"],
    )

    lp_lower_bounds[resource_var] = -resource["count"]

for column_id, column in lp_columns.items():
    to_add = defaultdict(float)
    for variable, coeff in column.items():
        if abs(coeff) < EPSILON:
            print(f"WARNING: zero or near-zero coeff: column_id={column_id} variable={variable} coeff={coeff}")

        if variable.startswith("item|") and coeff > 0:
            item_class = variable[5:]
            if item_class not in items:
                print(f"WARNING: item not found in items dict: {item_class}")
                continue

            item = items[item_class]
            form = item["form"]
            conveyance_limit = get_form_conveyance_limit(form)
            conveyance = coeff / conveyance_limit

            if column_id.startswith("miner|"):
                to_add["transport_power_cost"] += args.transport_power_cost * conveyance
                to_add["drone_battery_cost"] += args.drone_battery_cost * conveyance

            if form == "RF_SOLID":
                to_add["conveyors"] += conveyance
            else:
                to_add["pipelines"] += conveyance

    for variable, coeff in to_add.items():
        if coeff != 0.0:
            column[variable] = column.get(variable, 0.0) + coeff

for objective in ["points", "machines", "conveyors", "pipelines"]:
    column = {
        objective: -1,
    }
    add_lp_column(
        column,
        type_="objective",
        name=objective,
    )
    lp_equalities[objective] = 0.0

for extra_cost, cost_variable, cost_coeff in [
    ("transport_power_cost", "power_consumption", 1.0),
    ("drone_battery_cost", f"item|{BATTERY_CLASS}", -1.0),
]:
    column = {
        extra_cost: -1,
        cost_variable: cost_coeff,
    }
    add_lp_column(
        column,
        type_="extra_cost",
        name=extra_cost,
    )
    lp_equalities[extra_cost] = 0.0

column = {
    "power_consumption": -1,
    "power_production": -1,
}
add_lp_column(
    column,
    type_="power",
    name="usage",
)
lp_equalities["power_consumption"] = 0.0
lp_lower_bounds["power_production"] = 0.0

column = {
    "power_shard_usage": -1,
    "power_shards": -1,
}
add_lp_column(
    column,
    type_="objective",
    name="power_shards",
)
lp_equalities["power_shard_usage"] = 0.0
lp_lower_bounds["power_shards"] = -TOTAL_POWER_SHARDS

# pprint(lp_columns)
# pprint(lp_equalities)
# pprint(lp_lower_bounds)

def get_all_variables():
    variables = set()

    for column_id, column in lp_columns.items():
        for variable, coeff in column.items():
            variables.add(variable)

    for variable in variables:
        if variable not in lp_equalities and variable not in lp_lower_bounds:
            print(f"WARNING: no constraint for variable: {variable}")

    for variable in lp_equalities.keys():
        if variable not in variables:
            print(f"WARNING: equality constraint with unknown variable: {variable}")

    for variable in lp_lower_bounds.keys():
        if variable not in variables:
            print(f"WARNING: lower bound constraint with unknown variable: {variable}")

    return variables

variables = get_all_variables()
# pprint(variables)

### Pruning ###

reachable_items = set()
while True:
    any_added = False
    for column_id, column in lp_columns.items():
        eligible = True
        to_add = set()
        for variable, coeff in column.items():
            if variable.startswith("item|") and variable not in reachable_items:
                if coeff > 0:
                    to_add.add(variable)
                elif coeff < 0:
                    eligible = False
                    break
        if eligible and to_add:
            any_added = True
            reachable_items |= to_add
    if not any_added:
        break

unreachable_items = set(v for v in variables if v.startswith("item|")) - reachable_items

print("pruning unreachable items:")
pprint(unreachable_items)

columns_to_prune = list()
for column_id, column in lp_columns.items():
    for variable, coeff in column.items():
        if variable in unreachable_items and coeff < 0:
            columns_to_prune.append(column_id)
            break
for column_id in columns_to_prune:
    # pprint(lp_columns[column_id])
    del lp_columns[column_id]
for item_var in unreachable_items:
    if item_var in lp_equalities:
        del lp_equalities[item_var]

variables = get_all_variables()
# pprint(variables)

# pprint(lp_columns)
# pprint(lp_equalities)
# pprint(lp_lower_bounds)

### LP run ###

def to_index_map(seq):
    return {value: index for index, value in enumerate(seq)}

def from_index_map(d):
    result = [None] * len(d)
    for value, index in d.items():
        result[index] = value
    return result

# order is for report display, but we might as well sort it here
column_type_order = to_index_map(["objective", "power", "extra_cost", "sink", "waste", "manufacturer", "miner", "generator"])
column_subtype_order = to_index_map(["impure", "normal", "pure"])
objective_order = to_index_map(["points", "machines", "conveyors", "pipelines", "power_shards"])
extra_cost_order = to_index_map(["transport_power_cost", "drone_battery_cost"])

def column_order_key(arg):
    column_id, column = arg
    info = column.display_info

    type_ = info["type"]
    if type_ in column_type_order:
        type_key = (0, column_type_order[type_])
    else:
        type_key = (1, type_)

    name = info["display_name"]
    if type_ == "objective":
        name_key = objective_order[name]
    elif type_ == "extra_cost":
        name_key = extra_cost_order[name]
    else:
        name_key = name

    subtype = info["subtype"]
    if subtype in column_subtype_order:
        subtype_key = (0, column_subtype_order[subtype])
    else:
        subtype_key = (1, subtype)

    return (type_key, name_key, subtype_key, info["clock"], column_id)

sorted_columns = sorted(lp_columns.items(), key=column_order_key)
indices_eq = to_index_map(sorted(lp_equalities.keys()))
indices_lb = to_index_map(sorted(lp_lower_bounds.keys()))

# pprint(indices_eq)
# pprint(indices_lb)

lp_c = np.zeros(len(lp_columns), dtype=np.double)
lp_A_eq = np.zeros((len(lp_equalities), len(lp_columns)), dtype=np.double)
lp_b_eq = np.zeros(len(lp_equalities), dtype=np.double)
lp_A_lb = np.zeros((len(lp_lower_bounds), len(lp_columns)), dtype=np.double)
lp_b_lb = np.zeros(len(lp_lower_bounds), dtype=np.double)

objective_weights = {f"objective|{obj}": weight for (obj, weight) in {
    "points": 1,
    "machines": -args.machine_penalty,
    "conveyors": -args.conveyor_penalty,
    "pipelines": -args.pipeline_penalty,
    "power_shards": -args.power_shard_penalty_ratio * args.machine_penalty,
}.items()}

for column_index, (column_id, column) in enumerate(sorted_columns):
    if column_id in objective_weights:
        lp_c[column_index] = objective_weights[column_id]
    for variable, coeff in column.items():
        if variable in lp_equalities:
            lp_A_eq[indices_eq[variable], column_index] = coeff
        else:
            lp_A_lb[indices_lb[variable], column_index] = coeff

for variable, rhs in lp_equalities.items():
    lp_b_eq[indices_eq[variable]] = rhs

for variable, rhs in lp_lower_bounds.items():
    lp_b_lb[indices_lb[variable]] = rhs

print("running LP")

lp_result = scipy.optimize.linprog(-lp_c, A_ub=-lp_A_lb, b_ub=-lp_b_lb, A_eq=lp_A_eq, b_eq=lp_b_eq, method="highs")

if lp_result.status != 0:
    print("ERROR: LP did not terminate successfully")
    pprint(lp_result)
    sys.exit(1)

pprint(lp_result)

### Display formatting ###

def format_subtype(subtype):
    if subtype is None or subtype == "extractor":
        return None
    return re.sub(r"^BP_FrackingCore_?", "#", subtype).capitalize()

def get_column_desc(column):
    info = column.display_info
    tokens = [info["machine_name"] or info["type"], info["display_name"]]
    subtype = format_subtype(info["subtype"])
    if subtype is not None:
        tokens.append(subtype)
    if info["clock"] is not None:
        clock_percent = 100.0 * info["clock"]
        tokens.append(f"{clock_percent}%")
    return "|".join(tokens)

column_results = [
    (column_id, column, lp_result.x[column_index])
    for column_index, (column_id, column) in enumerate(sorted_columns)
]

if not args.show_unused:
    column_results = list(filter(lambda x: abs(x[2]) > EPSILON, column_results))

variable_breakdowns = {variable: {"production": [], "consumption": []} for variable in variables}

for column_id, column, column_coeff in column_results:
    column_desc = get_column_desc(column)

    print(f"{column_desc} = {column_coeff}")

    for variable, coeff in column.items():
        rate = column_coeff * coeff
        source = {
            "desc": column_desc,
            "count": column_coeff,
            "rate": abs(rate),
        }
        if abs(rate) < EPSILON:
            continue
        elif rate > 0:
            variable_breakdowns[variable]["production"].append(source)
        else:
            variable_breakdowns[variable]["consumption"].append(source)

variable_order = to_index_map(
    from_index_map(objective_order)
    + ["power_production", "power_consumption"]
    + from_index_map(extra_cost_order)
    + ["power_shards", "power_shard_usage", "item", "resource"]
)

def get_variable_display_info(variable):
    tokens = variable.split("|")
    type_ = tokens[0]
    if type_ == "item" or type_ == "resource":
        item_class = tokens[1]
        tokens[1] = get_item_display_name(item_class)
    return (type_, "|".join(tokens))

def finalize_variable_budget_side(budget_side):
    if not budget_side:
        return
    total_rate = 0.0
    for entry in budget_side:
        total_rate += entry["rate"]
    for entry in budget_side:
        entry["share"] = entry["rate"] / total_rate
    budget_side.sort(key=lambda entry: (-entry["share"], entry["desc"]))
    budget_side.insert(0, {"desc": "Total", "count": "n/a", "rate": total_rate, "share": 1.0})

for variable, breakdown in variable_breakdowns.items():
    type_, name = get_variable_display_info(variable)

    # don't show offsetting dummy items in the breakdown (e.g. "objective|points" as consumer of points)
    # currently these are precisely the consumption of special variables, but that may change
    if type_ not in ["item", "resource"]:
        breakdown["consumption"] = []

    breakdown["type_order"] = variable_order[type_]
    breakdown["name"] = name
    if variable in indices_lb:
        slack = lp_result.slack[indices_lb[variable]]
        if slack < -EPSILON:
            print(f"WARNING: lower bound violation: variable={variable} slack={slack}")
        breakdown["initial"] = -lp_lower_bounds[variable]
        breakdown["final"] = slack
    else:
        con = lp_result.con[indices_eq[variable]]
        if abs(con) > EPSILON:
            print(f"WARNING: equality constraint violation: variable={variable} con={con}")
    finalize_variable_budget_side(breakdown["production"])
    finalize_variable_budget_side(breakdown["consumption"])

sorted_variable_breakdowns = sorted(variable_breakdowns.values(), key=lambda bd: (bd["type_order"], bd["name"]))

# pprint(sorted_variable_breakdowns)

if args.xlsx_report:
    print("writing xlsx report")

    import xlsxwriter
    workbook = xlsxwriter.Workbook(args.xlsx_report)

    default_format = workbook.add_format({"align": "center"})
    top_format = workbook.add_format({"align": "center", "top": True})
    bold_format = workbook.add_format({"align": "center", "bold": True})
    bold_underline_format = workbook.add_format({"align": "center", "bold": True, "underline": True})
    bold_top_format = workbook.add_format({"align": "center", "bold": True, "top": True})
    bold_underline_top_format = workbook.add_format({"align": "center", "bold": True, "underline": True, "top": True})
    percent_format = workbook.add_format({"align": "center", "num_format": "0.0#####%"})

    sheet1 = workbook.add_worksheet("List" + args.xlsx_sheet_suffix)
    sheet2 = workbook.add_worksheet("Breakdown" + args.xlsx_sheet_suffix)

    def write_cell(sheet, *args, fmt=default_format):
        sheet.write(*args, fmt)

    sheet1.add_table(0, 0, len(column_results), 5, {
        "columns": [{"header": header, "header_format": bold_format}
        for header in ["Type", "Name", "Machine", "Subtype", "Clock", "Quantity"]],
        "style": "Table Style Light 16",
    })

    for i, (column_id, column, column_coeff) in enumerate(column_results):
        info = column.display_info
        write_cell(sheet1, i + 1, 0, info["type"])
        write_cell(sheet1, i + 1, 1, info["display_name"])
        write_cell(sheet1, i + 1, 2, info["machine_name"] or "n/a")
        write_cell(sheet1, i + 1, 3, info["subtype"] or "n/a")
        write_cell(sheet1, i + 1, 4, info["clock"] or "n/a", fmt=percent_format)
        write_cell(sheet1, i + 1, 5, column_coeff)

    for c, width in enumerate([14, 39, 25, 19, 11, 13]):
        sheet1.set_column(c, c, width)

    current_row = 0
    max_budget_entries = 0
    budget_rows = [
        ("desc", "Producer", "Consumer"),
        ("count", "Producer Count", "Consumer Count"),
        ("rate", "Production Rate", "Consumption Rate"),
        ("share", "Production Share", "Consumption Share"),
    ]

    production_share_cf = {
        "type": "2_color_scale",
        "min_type": "num",
        "max_type": "num",
        "min_value": 0,
        "max_value": 1,
        "min_color": "#FFFFFF",
        "max_color": "#99FF99"
    }
    consumption_share_cf = production_share_cf.copy()
    consumption_share_cf["max_color"] = "#FFCC66"

    for variable_index, breakdown in enumerate(sorted_variable_breakdowns):
        for budget_side_index, budget_side_name in enumerate(["production", "consumption"]):
            budget_side = breakdown[budget_side_name]
            if not budget_side:
                continue
            for budget_row in budget_rows:
                key = budget_row[0]
                name = budget_row[budget_side_index + 1]
                if key == "desc":
                    fmts = (bold_top_format, bold_underline_top_format)
                elif key == "share":
                    fmts = (bold_format, percent_format)
                else:
                    fmts = (bold_format, default_format)
                write_cell(sheet2, current_row, 0, breakdown["name"], fmt=fmts[0])
                write_cell(sheet2, current_row, 1, name, fmt=fmts[0])
                for i, entry in enumerate(budget_side):
                    write_cell(sheet2, current_row, i + 2, entry[key], fmt=fmts[1])
                if key == "share":
                    cf = production_share_cf if budget_side_name == "production" else consumption_share_cf
                    sheet2.conditional_format(current_row, 3, current_row, len(budget_side) + 1, cf)
                max_budget_entries = max(max_budget_entries, len(budget_side))
                current_row += 1
        for key in ["initial", "final"]:
            if key in breakdown:
                if key == "initial":
                    fmts = (bold_top_format, top_format)
                else:
                    fmts = (bold_format, default_format)
                fmt = bold_top_format if key == "initial" else bold_format
                write_cell(sheet2, current_row, 0, breakdown["name"], fmt=fmts[0])
                write_cell(sheet2, current_row, 1, key.capitalize(), fmt=fmts[0])
                write_cell(sheet2, current_row, 2, breakdown[key], fmt=fmts[1])
                current_row += 1

    for c, width in enumerate([41, 19, 13] + [59] * (max_budget_entries - 1)):
        sheet2.set_column(c, c, width)

    workbook.close()

reason: import numpy,import scipy