text:
from math import *
import random
import numpy as np

from optimize import OptimizationBasis, genetic_algorithm

def seq(start, end, step_count = None, step_size = None):
    assert step_count or step_size, "either step_count or step_size must be defined"
    step_size = step_size or (end-start) / step_count
    step_count = step_count or int(abs(end - start) / step_size)
    for i in range(0,step_count):
        yield start + i*step_size
def integrate(f, start, end, step_count = None, step_size = None):
    step_size = step_size or (end-start) / step_count
    step_count = step_count or int(abs(end - start) / step_size)
    return sum([f(x) for x in seq(start, end, step_count, step_size)]) * step_size

def rho(x,y,r): return exp( sqrt(x*x + y*y) - r )
def I(x,y,r):   return integrate(lambda xi: rho(xi,y,r), 0, x, step_count=100)
def F(x,y,r):   return rho(x,y,r)/I(x,y,r) - x/sqrt(x**2+y**2) - 1/x
def x0(x,r):    return sqrt(max(r*r - x*x,0))

SAMPLE_COUNT = 30
R = 10
CACHE = [
    ((x, y), F(x,y,R))
    for x in seq(-R,R,SAMPLE_COUNT)
    for y in seq(-x0(x,R), x0(x,R), SAMPLE_COUNT)
    if x != 0
]

numerator_basis = [
    '1',
    'x',
    #'y',
    #'x**2',
    #'x*y',
    'y**2',
    'x**3',
    #'x**2*y',
    'x*y**2',
    #'y**3',
    #'x**4',
    #'x**3*y',
    #'x**2*y**2',
    #'x*y**3',
    #'y**4',
    #'sqrt(x*x+y*y)',
    #'sqrt(x)',
    #'sqrt(y)',
]
denominator_basis = [
    '1',
    'x',
    # 'y',
    'x**2',
    #'x*y',
    'y**2',
    'x**3',
    #'x**2*y',
    'x*y**2',
    #'y**3',
    #'x**4',
    #'x**3*y',
    #'x**2*y**2',
    #'x*y**3',
    #'y**4',
    #'sqrt(x*x+y*y)',
    #'sqrt(x)',
    #'sqrt(y)',
]
def F2(input,params):
    (x,y) = input
    numerator_basis = [
        1,
        x,
        #y,
        #x**2,
        #x*y,
        y**2,
        x**3,
        #x**2*y,
        x*y**2,
        #y**3,
        #x**4,
        #x**3*y,
        #x**2*y**2,
        #x*y**3,
        #y**4,
        #sqrt(x*x+y*y),
        #sqrt(x),
        #sqrt(y),
    ]
    denominator_basis = [
        1,
        x,
        # y,
        x**2,
        #x*y,
        y**2,
        x**3,
        #x**2*y,
        x*y**2,
        #y**3,
        #x**4,
        #x**3*y,
        #x**2*y**2,
        #x*y**3,
        #y**4,
        #sqrt(x*x+y*y),
        #sqrt(x),
        #sqrt(y),
    ]
    a = params[0:len(numerator_basis)]
    b = params[len(numerator_basis):(len(numerator_basis)+len(denominator_basis))]
    return sum([a[i]*numerator_basis[i] for i in range(len(numerator_basis))]) /      \
           sum([exp(b[i])*denominator_basis[i] for i in range(len(denominator_basis))])

def code(params):
    a = params[0:len(numerator_basis)]
    b = params[len(numerator_basis):(len(numerator_basis)+len(denominator_basis))]
    numerator = '+'.join(f'{a[i]:.3f}*{numerator_basis[i]}' for i in range(len(a)))
    denominator = '+'.join(f'{exp(b[i]):.3f}*{denominator_basis[i]}' for i in range(len(b)))
    return f'({numerator}) / ({denominator})'
def pretty_basis(basis):
    return basis.replace('**3','³').replace('**2','²').replace('*','').replace('sqrt(','√').replace(')','')
def pretty(params):
    a = params[0:len(numerator_basis)]
    b = params[len(numerator_basis):(len(numerator_basis)+len(denominator_basis))]
    numerator = ' '.join(f'{a[i]:+.2f}{pretty_basis(numerator_basis[i])}' for i in range(len(a)))
    denominator = ' '.join(f'{exp(b[i]):+.2f}{pretty_basis(denominator_basis[i])}' for i in range(len(b)))
    return \
f"""({numerator})
{'-'*max(len(numerator),len(denominator))}
({denominator})"""

basis = OptimizationBasis(F2, pretty)

solutions = [np.array([random.gauss(0,1) for j in range(len(numerator_basis)+len(denominator_basis))]) for i in range(2000)]
solutions = genetic_algorithm(basis, CACHE, solutions)
print(code(solutions[0]))


text:
"""
    Main module for running NetSpeedGraphs.
"""

##### IMPORTS #####
# Standard imports
from pathlib import Path
from datetime import datetime, timedelta
from argparse import ArgumentParser

# Third party imports
import speedtest
import numpy as np
import pandas as pd
from bokeh.plotting import figure, output_file, save
from bokeh.models.formatters import DatetimeTickFormatter
from bokeh.models import (ColumnDataSource, DataTable, TableColumn,
                          NumberFormatter, DateFormatter)
from bokeh.layouts import grid

##### CONSTANTS #####
DATA_HEADER = ['Time', 'Ping (ms)', 'Download Speed (Mbs)',
               'Upload Speed (Mbs)']

##### FUNCTIONS #####
def allTests():
    """ Runs ping, download and upload speed tests.

    Returns
    -------
    ping: float
        Ping value in miliseconds.
    download: float
        Download speed in Mbs.
    upload: float
        Upload speed in Mbs.
    """
    st = speedtest.Speedtest()
    server = st.get_best_server()
    down = st.download()
    up = st.upload()
    return server['latency'], down / 1e6, up / 1e6

def plotGraph(data, path):
    """ Plots a graph with the download and upload speeds and pings.

    Parameters
    ----------
    data: pandas.DataFrame
        DataFrame containing 4 columns:
            - Time: datetime objects for the time the test was ran.
            - Ping: ping in milliseconds.
            - Download: download speed in megabits per second.
            - Upload: upload speed in megabits per second.
    path: pathlib Path
        Path to html file for outputting plots to.

    See Also
    --------
    readResults
    """
    # output to static HTML file
    output_file(path)

    # Use the pandas dataframe as the source
    source = ColumnDataSource(data)
    # Create a new plot and set x-axis type as datetime
    netPlot = figure(title="Network Speeds", x_axis_type='datetime',
                     x_axis_label='Time of Test',
                     y_axis_label='Speed (Mbs) / Ping (ms)',
                     tools=['xpan', 'xwheel_zoom', 'box_select', 'reset'],
                     active_drag='xpan', active_scroll='xwheel_zoom',
                     sizing_mode='stretch_both')
    # Change x axis tick format depending on zoom level
    netPlot.xaxis.formatter = DatetimeTickFormatter(hours = ['%H:%M'],
                                                    days = ['%d/%m/%Y'],
                                                    months = ['%m/%Y'])

    # Add the lines to the plot in different colours
    WIDTH = 2
    netPlot.line(x='Time', y='Ping', source=source, legend_label='Ping',
                 line_color='orange', line_width=WIDTH)
    netPlot.line(x='Time', y='Download', source=source, legend_label='Download',
                 line_color='blue', line_width=WIDTH)
    netPlot.line(x='Time', y='Upload', source=source, legend_label='Upload',
                 line_color='green', line_width=WIDTH)

    # Create table
    numFormatter = NumberFormatter(format='0.00')
    columns = [
        TableColumn(field="Time", title='Time',
                    formatter=DateFormatter(format="%Y-%m-%d %H:%M")),
        TableColumn(field='Ping', title='Ping (ms)', formatter=numFormatter),
        TableColumn(field="Download", title='Download Speed (Mbs)',
                    formatter=numFormatter),
        TableColumn(field='Upload', title='Upload Speed (Mbs)',
                    formatter=numFormatter)
    ]
    table = DataTable(source=source, columns=columns, width=400,
                      sizing_mode='stretch_height')

    # Add plot to grid layout
    layout = grid([netPlot, table], ncols=2)

    # show the results
    save(layout)
    return

def storeResults(results, path):
    """ Save the network speed results to CSV containing all results.

    Will create a CSV if it doesn't exist, or append results to it if it does.

    Parameters
    ----------
    results: list-like of floats
        The results from a single run of the network test in the following
        order: ping (miliseconds), download speed (Mbs) and upload speed (Mbs).
    path: pathlib Path
        Path to csv file for saving results to.

    See Also
    --------
    allTests
    """
    # Get current time of results
    now = datetime.now()
    # Create row of results
    row = [now.isoformat(), *[str(i) for i in results]]

    # Check if file exists and create it with header if not
    # then append current results to it
    header = not path.exists()
    with open(path, 'at') as out:
        if header:
            out.writelines(','.join(DATA_HEADER) + '\n')
        out.write(','.join(row) + '\n')
    return

def readResults(path):
    """ Read the csv containing all the results into a DataFrame.

    The `DATA_PATH` and `DATA_HEADER` constants are used when reading the csv.

    Parameters
    ----------
    path: pathlib Path
        Path to csv file for reading from.

    Returns
    -------
    data: pandas.DataFrame
        DataFrame containing 4 columns:
            - Time: datetime objects for the time the test was ran.
            - Ping: ping in milliseconds.
            - Download: download speed in megabits per second.
            - Upload: upload speed in megabits per second.
    """
    data = pd.read_csv(path, usecols=DATA_HEADER, parse_dates=[0])
    rename = {i: i.split()[0].strip().capitalize() for i in DATA_HEADER}
    data = data.rename(columns=rename)
    return data

def argParser():
    """ Creates an ArgumentParser to get output locations.

    Returns
    -------
    parser: argparse ArgumentParser
        Parser to get the output file locations from the arguments.
    """
    parser = ArgumentParser(description='Run a network test and update html plots.')
    parser.add_argument('data_file', type=Path,
                        help='csv file for storing all network test results.')
    parser.add_argument('html_file', type=Path,
                        help='html file for saving the output plots to.')
    return parser

def main():
    """ Runs the network test to get results then updates the csv and graphs. """
    # Get file locations from arguments
    parser = argParser()
    args = parser.parse_args()

    # Run a test and update the graphs
    netRes = allTests()
    storeResults(netRes, args.data_file)
    results = readResults(args.data_file)
    plotGraph(results, args.html_file)
    return

##### MAIN #####
if __name__ == '__main__':
    main()

ext: py

reason: import numpy


text:
import numpy as np
import tensorflow as tf

def project_in_2D(K, camera_pose, mesh, resolution_px):
    """
    Project all 3D triangle vertices in the mesh into
    the 2D image of given resolution

    Parameters
    ----------
    K: ndarray
        Camera intrinsics matrix, 3x3
    camera_pose: ndarray
        Camera pose (inverse of extrinsics), 4x4
    mesh: ndarray
        Triangles to be projected in 2d, (Nx3x3)
    resolution_px: tuple
        Resolution of image in pixel

    Returns
    -------
    coords_projected_2D: ndarray
        Triangle vertices projected in 2D and clipped to
        image resolution
    """
    resolution_x_px, resolution_y_px = resolution_px  # image resolution in pixels

    # Decompose camera pose into rotation and translation
    RT = camera_pose[:-1, :]  # remove homogeneous row
    R = RT[:, :-1]  # rotation matrix 3x3
    T = RT[:, -1:]  # translation vector 3x1

    # Invert the camera pose matrix to get the camera extrinsics
    # Due to the particular matrix geometry we can avoid raw inversion
    Rc = tf.matrix_transpose(R)
    Tc = tf.matmul(-Rc, T)
    RT = tf.concat([Rc, Tc], axis=-1)  # camera extrinsics

    # Correct reference system of extrinsics matrix
    #   y is down: (to align to the actual pixel coordinates used in digital images)
    #   right-handed: positive z look-at direction
    correction_factor = tf.constant(value=np.array([[1., 0., 0.],
                                                    [0., -1., 0.],
                                                    [0., 0., -1.]]), dtype=tf.float32)
    RT = tf.matmul(correction_factor, RT)

    # Compose whole camera projection matrix (3x4)
    P = tf.matmul(K, RT)

    mesh_flat = tf.reshape(mesh, shape=(-1, 3))
    len_mesh_flat = tf.shape(mesh_flat)[0]
    # Create constant tensor to store 3D model coordinates
    coords_3d_h = tf.concat([mesh_flat, tf.ones(shape=(len_mesh_flat, 1), dtype=tf.float32)], axis=-1)  # n_triangles, 4
    coords_3d_h = tf.transpose(coords_3d_h, perm=[1, 0])  # 4, n_triangles

    # Project 3D vertices into 2D
    coords_projected_2D_h = tf.transpose(tf.matmul(P, coords_3d_h), perm=[1, 0])  # n_triangles, 3
    coords_projected_2D = coords_projected_2D_h[:, :2] / (coords_projected_2D_h[:, 2:3] + 1e-8)

    # Clip indexes in image range
    coords_projected_2D_x_clip = tf.clip_by_value(coords_projected_2D[:, 0:0 + 1],
                                                  clip_value_min=-1, clip_value_max=resolution_x_px)
    coords_projected_2D_y_clip = tf.clip_by_value(coords_projected_2D[:, 1:1 + 1],
                                                  clip_value_min=-1, clip_value_max=resolution_y_px)

    return tf.concat([coords_projected_2D_x_clip, coords_projected_2D_y_clip], axis=-1)

def calibration_matrix(resolution_px, resolution_mm, focal_len_mm, skew=0.):
    """
    Return calibration matrix K given camera information
    """
    # Camera intrinsics parameters
    resolution_x_px, resolution_y_px = resolution_px  # image resolution in pixels
    resolution_x_mm, resolution_y_mm = resolution_mm  # size of camera sensor in mm

    skew = skew  # "skew param will be zero for most normal cameras" Hartley, Zisserman
    focal_len_mm = focal_len_mm  # camera focal len in mm

    assert (resolution_x_px / resolution_y_px == resolution_x_mm / resolution_y_mm)

    m_x = resolution_x_px / resolution_x_mm
    m_y = resolution_y_px / resolution_y_mm

    alpha_x = focal_len_mm * m_x  # focal length of the camera in pixels
    alpha_y = focal_len_mm * m_y  # focal length of the camera in pixels

    x_0 = resolution_x_px / 2
    y_0 = resolution_y_px / 2

    return np.array([[alpha_x, skew, x_0],
                     [0, alpha_y, y_0],
                     [0, 0, 1]])


text:
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable

class A3C_LSTM_GA(torch.nn.Module):

    def __init__(self, args):
        super(A3C_LSTM_GA, self).__init__()

        # General
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

        # Image Processing
        self.conv1 = nn.Conv2d(3, 128, kernel_size=8, stride=4)
        self.conv2 = nn.Conv2d(128, 64, kernel_size=4, stride=2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=4, stride=2)

        # Instruction Processing
        self.gru_hidden_size = 256
        self.input_size = args.input_size
        self.embedding = nn.Embedding(self.input_size, 32)
        self.gru = nn.GRU(32, self.gru_hidden_size)

        # Gated-Attention layers
        self.attn_linear = nn.Linear(self.gru_hidden_size, 64)

        # Time embedding layer, helps in stabilizing value prediction.
        # max episode length is 30 by default.
        self.time_emb_dim = 32
        self.time_emb_layer = nn.Embedding(
                300 + 1,  #This is because we train with a higher max episode length
                #args.max_episode_length+1,
                self.time_emb_dim)

        # Action embedding
        self.action_emb_dim = 8
        self.action_emb_layer = nn.Embedding(3, self.action_emb_dim) #Map from action to embedding

        self.frame_size = 64 * 8 * 17
        self.cnn_out_size = 5 * self.frame_size
        self.rep_size = 2 * self.cnn_out_size + self.gru_hidden_size
        # A3C-LSTM layers
        self.linear = nn.Linear(self.rep_size, 384)
        self.lstm = nn.LSTMCell(384, 256)
        self.critic_linear = nn.Linear(256 + self.time_emb_dim, 1)
        self.actor_linear = nn.Linear(256 + self.time_emb_dim, 3)

        #Inverse dynamic model
        self.inverse_linear = nn.Linear(2 * self.cnn_out_size, 256)
        self.inverse_actor = nn.Linear(256, 3)

        #Forward dynamic model. Action is embedded in 8-dim vector. Model predicts next final FRAME, not next 5 frames
        self.forward_linear = nn.Linear(self.cnn_out_size + self.action_emb_dim, 256)
        self.forward_state = nn.Linear(256, self.frame_size)

        self.train()

    #Produces action for state s_t
    def teacherForward(self, inputs):
        x, input_inst, (tx, hx, cx) = inputs

        # Get the image representation
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x_image_rep = self.relu(self.conv3(x))

        # Get the instruction representation
        embedded = self.embedding(input_inst).permute(1,0,2)
        encoder_hidden = Variable(self.gru(embedded)[0][-1]) #Get last hidden state
        x_instr_rep = encoder_hidden

        # Get the attention vector from the instruction representation
        x_attention = self.sigmoid(self.attn_linear(x_instr_rep))

        # Gated-Attention
        x_attention = x_attention.unsqueeze(2).unsqueeze(3)
        x_attention = x_attention.expand(5, 64, 8, 17)
        print(x_attention.shape)
        print(x_image_rep.size())
        assert x_image_rep.size() == x_attention.size()
        x = x_image_rep*x_attention
        #Concatentation as suggested by Manning
        x = torch.cat((x.view(-1), x_image_rep.view(-1), x_instr_rep.view(-1))).unsqueeze(0)

        # A3C-LSTM
        x = self.relu(self.linear(x))
        hx, cx = self.lstm(x, (hx, cx))
        time_emb = self.time_emb_layer(tx)
        x = torch.cat((hx, time_emb.view(-1, self.time_emb_dim)), 1)

        return self.critic_linear(x), self.actor_linear(x), (hx, cx)

    '''
    Predicts a_t based on s_t and s_{t+1}
    '''
    def inverseDynamicsForward(self, inputs):
        s1, s2 = inputs

        # Get the image representation for s1
        s1 = self.relu(self.conv1(s1))
        s1 = self.relu(self.conv2(s1))
        s1_image_rep = self.relu(self.conv3(s1)).view(-1)

        # Get the image representation for s2
        s2 = self.relu(self.conv1(s2))
        s2 = self.relu(self.conv2(s2))
        s2_image_rep = self.relu(self.conv3(s2)).view(-1)

        x = torch.cat((s1_image_rep, s2_image_rep))
        x = self.inverse_linear(x)
        x = self.inverse_actor(x)

        return x #this is logits over action. Embedding is just to process action.

    '''
    Predicts s_{t+1} given s_t and a_t. 0 < a_t < #actions
    Both s_t and a_t are variables
    #actions is 3 for now
    '''
    def forwardDynamicsForward(self, inputs):
        s1, a1 = inputs

        # Get the image representation for s1
        s1 = self.relu(self.conv1(s1))
        s1 = self.relu(self.conv2(s1))
        s1_image_rep = self.relu(self.conv3(s1)).view(-1)

        a1 = self.action_emb_layer(a1).view(-1)
        x = torch.cat((s1_image_rep, a1))

        x = self.forward_linear(x)
        x = self.forward_state(x)

        return x

    def getImageRep(self, s):
        s = self.relu(self.conv1(s))
        s = self.relu(self.conv2(s))
        s_image_rep = self.relu(self.conv3(s)).view(-1)
        return s_image_rep

    #Similar to teacherForward, but sets attention and instruction representation to 0.
    def curiousForward(self, inputs):
        x, _, (tx, hx, cx) = inputs

        # Get the image representation
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x_image_rep = self.relu(self.conv3(x))

        # Attention and instruction representation start out as 0
        x_instr_rep = torch.zeros((self.gru_hidden_size))
        x = torch.zeros((self.cnn_out_size))
        x = torch.cat((x.view(-1), x_image_rep.view(-1), x_instr_rep.view(-1))).unsqueeze(0)

        # A3C-LSTM
        x = self.relu(self.linear(x))
        hx, cx = self.lstm(x, (hx, cx))
        time_emb = self.time_emb_layer(tx)
        x = torch.cat((hx, time_emb.view(-1, self.time_emb_dim)), 1)

        return self.critic_linear(x), self.actor_linear(x), (hx, cx)

    def forward(self, inputs, teacher=True, inverse=True, curious=True):
        if teacher:
            return self.teacherForward(inputs)
        elif inverse:
            return self.inverseDynamicsForward(inputs)
        elif not curious:
            return self.forwardDynamicsForward(inputs)
        else:
            return self.curiousForward(inputs)


text:
import numpy as np
from ANN.Layer import *

class LeNet5:
    def __init__(self):
        self.conv1 = Conv2d(1, 6, 5, stride=1, padding=2)
        self.pooling1 = MaxPooling(kernel_size=2, stride=2)
        self.BN1 = BN(6, momentum=0.9, is_train=True)
        self.relu1 = Relu()

        self.conv2 = Conv2d(6, 16, 5, stride=1, padding=0)
        self.pooling2 = MaxPooling(kernel_size=2, stride=2)
        self.BN2 = BN(16, momentum=0.9, is_train=True)
        self.relu2 = Relu()

        self.conv3 = Conv2d(16, 120, 5, stride=1, padding=0)

        self.fc4 = Fc(120*1*1, 84)
        self.relu4 = Relu()
        self.fc5 = Fc(84, 10)

        self.softmax = Softmax_CrossEntropy()

    def forward(self, input, labels, is_train=True):
        x = self.conv1.forward(input)
        x = self.pooling1.forward(x)
        x = self.BN1.forward(x, is_train)
        x = self.relu1.forward(x)

        x = self.conv2.forward(x)
        x = self.pooling2.forward(x)
        x = self.BN2.forward(x, is_train)
        x = self.relu2.forward(x)

        x = self.conv3.forward(x)

        x = self.fc4.forward(x)
        x = self.relu4.forward(x)
        x = self.fc5.forward(x)

        x = self.softmax.forward(x)
        loss = self.softmax.calculate_loss(x, labels)

        return loss, x

    def backward(self, learning_rate):
        eta = self.softmax.gradient()

        eta = self.fc5.backward(eta, learning_rate)
        eta = self.relu4.backward(eta)
        eta = self.fc4.backward(eta, learning_rate)

        eta = self.conv3.backward(eta, learning_rate)

        eta = self.relu2.backward(eta)
        eta = self.BN2.backward(eta, learning_rate)
        eta = self.pooling2.backward(eta)
        eta = self.conv2.backward(eta, learning_rate)

        eta = self.relu1.backward(eta)
        eta = self.BN1.backward(eta, learning_rate)
        eta = self.pooling1.backward(eta)
        eta = self.conv1.backward(eta, learning_rate)

max_stars_repo_name: FFFjx/Deep-Learning

max_issues_repo_name: FFFjx/Deep-Learning

max_forks_repo_name: FFFjx/Deep-Learning


text:
from pathlib import Path

import numpy as np

from kAI.neural_network import *
from kAI.utils import row_vector_list, plot_training, save_stats
from mnist import mnist

if __name__ == '__main__':
    np.random.seed(100)
    # Define Network
    network = Network(
        [Layer(784), Layer(6, Sigmoid), Layer(6, Sigmoid), Layer(10, Sigmoid)])  # , ImpactRealtimeVizInter)

    # Load MNIST dataset
    classes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    p, t, test, labels = mnist()

    # Reshape
    p = row_vector_list(p, 784)
    t = row_vector_list(t, 10)
    test = row_vector_list(test, 784)
    labels = row_vector_list(labels, 10)

    # Define Optimizer
    o = GrowingStochasticGD(
        learning_rate=0.2,
        cost_func=SumSquaredError,
        patterns=p,
        targets=t,
        epochs=80,
        splitting_coefficient=1,
        test_acc=(test, labels),
        observer=True
    )
    o.epoch_skip = 1
    network.set_optimizer(o)

    # Start optimizing
    stats = (None, None, None, None, None)
    try:
        stats = network.optimize()
    except KeyboardInterrupt:
        pass
    mp = network.save_as(Path('models/'))
    save_stats(mp, *stats)

    plot_training(*stats, font_size=15)

max_stars_repo_name: theswampire/matura

max_issues_repo_name: theswampire/matura

max_forks_repo_name: theswampire/matura


text:
from __future__ import division
import matplotlib.pyplot as plt
import cPickle as pickle
import numpy as np
import os
import sys
import re

"""
File data input example:
#  PDB code |Reported in PDB  | Calc from NCS   |  ASU initial    |   ASU final     | Res.   |   NCS   |  Solvent |     Data     |  D/A ratio   |  D/A ratio   | Year   | Use  | Geo. | Trans. | time (sec)
#           | r-work | r-free | r-work | r-free | r-work | r-free | r-work | r-free |        | copies  | fraction | completeness |     ASU      |     NCS      |        | NCS  | Rest | refine |
#----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
     1vcr   | 0.3790 | 0.3530 | 0.4292 | 0.4486 | 0.4242 | 0.4382 | 0.3280 | 0.4407 |  9.50  |    5    |   0.93   |     1.00     |     0.21     |     1.06     |  2004  |False | True | False  |  1131

Using pdb file from local machine (not from MIRROR)
Using cif file from local machine (not from MIRROR)
#  PDB code |Reported in PDB  | Calc from NCS   |  ASU initial    |   ASU final     | Res.   |   NCS   |  Solvent |     Data     |  D/A ratio   |  D/A ratio   | Year   | Use  | Geo. | Trans. | time (sec)
#           | r-work | r-free | r-work | r-free | r-work | r-free | r-work | r-free |        | copies  | fraction | completeness |     ASU      |     NCS      |        | NCS  | Rest | refine |
#----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
     1vcr   | 0.3790 | 0.3530 | 0.4292 | 0.4486 | 0.4242 | 0.4382 | 0.3662 | 0.4119 |  9.50  |    5    |   0.93   |     1.00     |     0.21     |     1.06     |  2004  | True | True | False  |  928

Using pdb file from local machine (not from MIRROR)
Using cif file from local machine (not from MIRROR)
#  PDB code |Reported in PDB  | Calc from NCS   |  ASU initial    |   ASU final     | Res.   |   NCS   |  Solvent |     Data     |  D/A ratio   |  D/A ratio   | Year   | Use  | Geo. | Trans. | time (sec)
#           | r-work | r-free | r-work | r-free | r-work | r-free | r-work | r-free |        | copies  | fraction | completeness |     ASU      |     NCS      |        | NCS  | Rest | refine |
#----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
     1vcr   | 0.3790 | 0.3530 | 0.4292 | 0.4486 | 0.4242 | 0.4382 | 0.3558 | 0.4175 |  9.50  |    5    |   0.93   |     1.00     |     0.21     |     1.06     |  2004  | True | True |  True  |  1554
#----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# Rotation,Translation difference      |  Rxx   |  Rxy   |  Rxz   |  Ryx   |  Ryy   |  Ryz   |  Rzx   |  Rzy   |  Rzz   |   Tx   |   Ty   |   Tz
#----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
  Transform number: 1                  |-0.0132 | 0.0194 | 0.0100 |-0.0370 | 0.0708 | 0.0253 |-0.0481 |-0.0465 | 0.0563 | 0.8755 | 1.6418 | 3.8338
  Transform number: 2                  | 0.0307 |-0.0333 | 0.0440 |-0.0088 |-0.0076 |-0.0097 | 0.0241 | 0.0532 | 0.0213 |-2.8667 | 0.5622 |-3.2518
  Transform number: 3                  |-0.0204 | 0.0023 |-0.0292 | 0.0158 | 0.0264 |-0.0083 |-0.0071 |-0.0135 | 0.0055 | 2.3881 |-1.3371 | 0.1937
  Transform number: 4                  |-0.0186 | 0.0733 | 0.0095 |-0.0460 | 0.0168 | 0.0087 | 0.0651 | 0.0341 |-0.0072 |-0.4086 | 5.3377 |-6.0357

Data to excel output
--------------------
   0 (0)|  1 (10)        | 2(15)|   3(9)     |      4(12)       |   5(11)          |
pdb code| num ncs copies | year | resolution | data completeness| solvent fraction |

 6(13)         |   7(14)
data/param ASU | data/param NCS |

 8(1)             |     9(2)          |     10(5)      | 11(6)
R-work pdb header | R-free pdb header | R-work init    | init

  12(7)(16)               |        13(8)(16)          |
R-work without strict ncs | R-free without strict ncs |

    14(7)(16)             |    15(8)(16)
 R-work using strict ncs | R-free using strict ncs |

      16(7)(18)        |  17(8)(18)           |
 R-work with transform |R-free with transform |

    18(17)
 use geometry restraints |

 19(19)(16)   | 18(19)(16) |    20(19)(18)  |
time no ncs   | time ncs   | time_transform |

"""

class results_collection(object):

  def __init__(self):
    self.data_records = []
    self.data_records_long = []
    self.data_records_strict_ncs_dict = {}
    self.data_records_transform_dict = {}
    self.data_records_without_strict_ncs_dict = {}
    self.data_records_dict = {}
    self.transform_records_dict = {}
    self.cols_inp = [
      'pdb_code',
      'r_work_pdb_header',
      'r_free_pdb_header',
      'r_work_single_ncs',
      'r_free_single_ncs',
      'r_work_asu_init',
      'r_free_asu_init',
      'r_work_asu_final',
      'r_free_asu_final',
      'resolution',
      'num_ncs_copies',
      'solvent_fraction',
      'data_completeness',
      'data_to_atoms_ratio_asu',
      'data_to_atoms_ratio_ncs',
      'year',
      'use_ncs_constraints',
      'use_geometry_restraints',
      'use_transforms',
      'time']

    self.cols_names = [
      'pdb_code',
      'num_ncs_copies',
      'year',
      'resolution',
      'data_completeness',
      'solvent_fraction',
      'data_to_atoms_ratio_asu',
      'data_to_atoms_ratio_ncs',
      'r_work_pdb_header',
      'r_free_pdb_header',
      'r_work_asu_init',
      'r_free_asu_init',
      'r_work_asu_no_ncs',
      'r_free_asu_no_ncs',
      'r_work_asu_ncs',
      'r_free_asu_ncs',
      'r_work_asu_transform',
      'r_free_asu_transform',
      'use_geometry_restraints',
      'time_no_ncs',
      'time_ncs',
      'time_transform']

    self.cols_names_for_table1 = [x.replace('_',' ') for x in self.cols_inp]
    self.cols_names_for_table2 = [x.replace('_',' ') for x in self.cols_names]
    # map the location in cols_inp to cols_names
    self.map_to_no_ncs = self.get_map(use_ncs_constraints=False)
    self.map_to_ncs = self.get_map(use_ncs_constraints=True,use_transforms=False)
    self.map_to_transform=self.get_map(use_ncs_constraints=True,use_transforms=True)
    # set records type
    self.set_records_type()

    class sort_inp():
      def __init__(self):
        """ Parameter number to sort by """
    sort_inp.__dict__.update(zip(self.cols_inp, range(len(self.cols_inp))))
    self.sort_inp = sort_inp

    class sort_out():
      def __init__(self):
        """ Parameter number to sort by """
    sort_out.__dict__.update(zip(self.cols_names, range(len(self.cols_names))))
    self.sort_out = sort_out

  def set_records_type(self):
    """ Create list containing the type of each record """
    strings = ['pdb_code']
    floats = [
      'r_work_pdb_header',
      'r_free_pdb_header',
      'r_work_single_ncs',
      'r_free_single_ncs',
      'r_work_asu_init',
      'r_free_asu_init',
      'r_work_asu_final',
      'r_free_asu_final',
      'resolution',
      'solvent_fraction',
      'data_completeness',
      'data_to_atoms_ratio_asu',
      'data_to_atoms_ratio_ncs',
      'time']
    integers = ['num_ncs_copies','year']
    bools = ['use_ncs_constraints', 'use_geometry_restraints', 'use_transforms']

    r = range(len(self.cols_inp))
    self.float_type_records = [i for i in r if self.cols_inp[i] in floats]
    self.int_type_records = [i for i in r if self.cols_inp[i] in integers]
    self.bool_type_records = [i for i in r if self.cols_inp[i] in bools]
    self.str_type_records = [i for i in r if self.cols_inp[i] in strings]

  def get_map(self,use_ncs_constraints=True,use_transforms=False):
    """
    set the map between the input records and the combined output line

    Args:
      use_ncs_constraints (bool): when True, using strict NCS
      use_transforms (bool): when True, refining rotation and translation

    Returns:
      (dict): map record location in the input to the record number at the
        output
    """
    if not use_ncs_constraints:
      d = 0
    elif use_transforms:
      d = 2
    else:
      d = 1

    direct_match = ['pdb_code','r_work_pdb_header','r_free_pdb_header',
                    'r_work_asu_init','r_free_asu_init','resolution',
                    'num_ncs_copies','solvent_fraction','data_completeness',
                    'data_to_atoms_ratio_asu','data_to_atoms_ratio_ncs','year',
                    'use_geometry_restraints']
    type_match = [['time',['time_no_ncs','time_ncs','time_transform']],
                  ['r_work_asu_final',['r_work_asu_no_ncs','r_work_asu_ncs',
                                       'r_work_asu_transform',]],
                  ['r_free_asu_final',['r_free_asu_no_ncs','r_free_asu_ncs',
                                       'r_free_asu_transform',]]]
    map_dict = {}
    r1 = range(len(self.cols_inp))
    r2 = range(len(self.cols_names))
    for val in direct_match:
      k = [i for i in r1 if val == self.cols_inp[i]][0]
      v = [i for i in r2 if val == self.cols_names[i]][0]
      map_dict[k] = v

    for v1,v_all in type_match:
      v2 = v_all[d]
      k = [i for i in r1 if v1 == self.cols_inp[i]][0]
      v = [i for i in r2 if v2 == self.cols_names[i]][0]
      map_dict[k] = v

    return map_dict

  def read_filenames(self):
    """
    """
    use_rec = self.sort_inp
    # Read files in current directory
    files = os.listdir(os.getcwd())
    # collect only non-empty files that starts with log_
    files = [x for x in files if x.startswith('log_') and os.stat(x).st_size>0]
    print 'Number of log files with data: {}'.format(len(files))
    # prepare regular expression search of all lines that start with pdb code
    regex = re.compile('     ....   \|')
    for fn in files:
      pdb_code = fn[4:]
      d = open(fn, "r").readlines()
      for ln in d:
        # print info and warnings
        print_warning(ln,pdb_code)
        # process files with good data (line start with pdb code)
        if re.match(regex,ln[:13]):
          data = self.process_data_according_to_data_type(ln)
          # collect data as is
          self.data_records.append(data)
          # combine all data for each pdb file to one row of data
          data_record = self.combine_records(data)
          self.data_records_dict[data[0]] = data_record
          # build other dictionaries
          if data[use_rec.use_ncs_constraints]:
            self.data_records_strict_ncs_dict[data[0]] = data
          else:
            self.data_records_without_strict_ncs_dict[data[0]] = data
          # build other dictionaries
          if data[use_rec.use_transforms]:
            self.data_records_transform_dict[data[0]] = data
        # process transformation info
        if ln.startswith('  Transform number:'):
          transform_data = [x.strip() for x in ln.split('|')][1:]
          transform_data = map(float,transform_data)
          if self.transform_records_dict.has_key(data[0]):
            self.transform_records_dict[data[0]].append(transform_data)
          else:
            self.transform_records_dict[data[0]] = [transform_data]

    # create a list version of the dictionary
    for key,val in self.data_records_dict.iteritems():
      self.data_records_long.append(val)
    print '-'*60
    print 'There are {} good data records'.format(len(self.data_records))
    print 'There are {} pdb records'.format(len(self.data_records_long))
    print '-'*60

    # organize data
    # Sort records
    # self.data_records.sort(key=lambda x:x[use_rec.time])
    print 'done with collection'
    pickle.dump(self.data_records,open('data_records_v2','w'))
    pickle.dump(self.data_records_long,open('data_records_long_v2','w'))
    pickle.dump(self.data_records_dict,open('data_records_dict_v2','w'))
    pickle.dump(self.transform_records_dict,open('transform_records_dict','w'))
    pickle.dump(
      self.data_records_strict_ncs_dict,open('data_records_strict_ncs_dict_v2','w'))
    pickle.dump(
      self.data_records_without_strict_ncs_dict,
      open('data_records_without_strict_ncs_dict_v2','w'))

  def get_data_from_files(self):
    """  Read data from files    """
    self.data_records = pickle.load(open('data_records_v2','r'))
    self.data_records_long = pickle.load(open('data_records_long_v2','r'))
    self.data_records_dict = pickle.load(open('data_records_dict_v2','r'))
    self.transform_records_dict =pickle.load(open('transform_records_dict','r'))
    self.data_records_strict_ncs_dict = pickle.load(
      open('data_records_strict_ncs_dict_v2','r'))
    self.data_records_without_strict_ncs_dict = pickle.load(
      open('data_records_without_strict_ncs_dict_v2','r'))
    print 'Got the data from files'

  def plot_results_1(self):
    """
    Plot a result summery:
    r-work (refined using strict-ncs) vs.
    r-work (refined without strict-ncs)

    ncs_err , no_ncs_err are the difference between r-work and r-free for the two
    refinement methods

    The size of the circles indicate the size of (ncs_err - no_ncs_err). It is blue
    when the strict_ncs refinement make the (r_work - r_free) smaller.
    """
    rec_num = self.sort_inp
    # Collect data points that have results with and without strict_ncs
    x_ncs = []; y_no_ncs =[]
    x2 = []; y2 = []
    time_ncs = []; time_no_ncs = []
    pdb_code = []
    for x in self.data_records_strict_ncs_dict:
      if self.data_records_without_strict_ncs_dict.has_key(x):
        x_ncs.append(self.data_records_strict_ncs_dict[x]
                           [rec_num.r_work_asu_final])
        y_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                                  [rec_num.r_work_asu_final])
        x2.append(self.data_records_strict_ncs_dict[x]
                      [rec_num.r_free_asu_final])
        y2.append(self.data_records_without_strict_ncs_dict[x]
                      [rec_num.r_free_asu_final])
        time_ncs.append(self.data_records_strict_ncs_dict[x]
                      [rec_num.time])
        time_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                      [rec_num.time])
        pdb_code.append(x)

    # variable error bar values
    ncs_err = [abs(x-y) for (x,y) in zip(x_ncs,x2)]
    no_ncs_err = [abs(x-y) for (x,y) in zip(y_no_ncs,y2)]

    # get larges r-values for the 45 degrees line
    maxval = max(x_ncs + y_no_ncs) *1.05

    # Convert lists to numpy arrays
    x_ncs = np.array(x_ncs)
    y_no_ncs = np.array(y_no_ncs)
    ncs_err = np.array(ncs_err)
    no_ncs_err = np.array(no_ncs_err)

    # First illustrate basic pyplot interface, using defaults where possible.
    plt.figure()
    colors = []
    # delta_errs = [(x-y) for (x,y) in zip (x_ncs,y_no_ncs)]
    delta_errs = [(x-y) for (x,y) in zip (ncs_err,no_ncs_err)]
    # delta_errs_percent=[(x-y)/r for (x,y,r) in zip (ncs_err,no_ncs_err,x_ncs)]
    point_size = lambda x: 5 + 1000 * abs(x)
    s = [point_size(x) for x in delta_errs]
    for i,delta_err in enumerate(delta_errs):
      if delta_err > 0:
        colors.append('y')
        print 'Smaller (r_work - r_free) without NCS: ',pdb_code[i]
      else:
        colors.append('b')

    # add points for size reference : Difference of 0, 0.05, 0.1
    s.extend([point_size(0.1),point_size(0.05),point_size(0)])
    colors.extend(['g','g','g'])
    x_ncs = np.append(x_ncs,[.02,.02,.02])
    d = maxval - 0.55
    y_no_ncs = np.append(y_no_ncs,[.4 + d,.45 + d,.5 + d])
    plt.scatter(x_ncs,y_no_ncs,s,c=colors)
    plt.plot([0,maxval],[0,maxval])
    plt.xlabel('r-work Refinement using strict-ncs')
    plt.ylabel('r-work Refinement without strict-ncs')
    plt.title('strict-ncs influence on (R_work - R_free)')
    # plot reference point, to indicate the meaning of size
    plt.text(0.035,0.49 + d, 'The same r_work - r_free value',fontsize=14)
    plt.text(0.035,0.44 + d, '0.05 difference',fontsize=14)
    plt.text(0.035,0.39 + d, '0.1 difference',fontsize=14)

    plt.xlim(0,maxval)
    plt.ylim(0,maxval)
    plt.show()

  def plot_results_2(self):
    """
    Plot a result summery:
    time for refinement with strict ncs vs. without
    """
     # Collect data points that have results with and without strict_ncs
    time_ncs = []; time_no_ncs = []
    pdb_code = []
    for x in self.data_records_strict_ncs_dict:
      if self.data_records_without_strict_ncs_dict.has_key(x):
        time_ncs.append(self.data_records_strict_ncs_dict[x]
                      [self.sort_inp.time])
        time_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                      [self.sort_inp.time])
        pdb_code.append(x)

    # print time outliers
    # d = 0.1
    # for (x,y,code) in zip(time_ncs,time_no_ncs,pdb_code):
    #   if abs(x-y)/x > d:
    #     print 'Time difference is {0:.0f}% for {1}'.format(100*(x-y)/x,code)

    # get larges r-values for the 45 degrees line
    maxval = max(time_ncs + time_no_ncs) *1.05

    # Convert lists to numpy arrays
    x_ncs = np.array(time_ncs)
    y_no_ncs = np.array(time_no_ncs)
    # plot
    plt.figure()
    plt.plot(x_ncs,y_no_ncs,'o',[0,maxval],[0,maxval])
    plt.xlabel('Time[sec] Refinement using strict-ncs')
    plt.ylabel('Time[sec] Refinement without strict-ncs')
    plt.title('Looking at refinement time when using stric_ncs')
    plt.xlim(0,maxval)
    plt.ylim(0,maxval)
    plt.show()

  def plot_results_3(self):
    """
    Plot a result summery:
    r-work (refined using strict-ncs and Transforms) vs.
    r-work (refined without strict-ncs)

    ncs_err , no_ncs_err are the difference between r-work and r-free for the two
    refinement methods

    The size of the circles indicate the size of (ncs_err - no_ncs_err). It is blue
    when the strict_ncs refinement make the (r_work - r_free) smaller.
    """
    sort_by = self.sort_inp
    # Collect data points that have results with and without strict_ncs
    x_ncs = []; y_no_ncs =[]
    x2 = []; y2 = []
    pdb_code = []
    print len(self.data_records_transform_dict)
    for x in self.data_records_transform_dict:
      if self.data_records_without_strict_ncs_dict.has_key(x):
        x_ncs.append(self.data_records_transform_dict[x]
                           [sort_by.r_work_asu_final])
        y_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                                  [sort_by.r_work_asu_final])
        x2.append(self.data_records_transform_dict[x]
                      [sort_by.r_free_asu_final])
        y2.append(self.data_records_without_strict_ncs_dict[x]
                      [sort_by.r_free_asu_final])
        pdb_code.append(x)

    # get larges r-values for the 45 degrees line
    maxval = max(x_ncs + y_no_ncs) *1.05

    # variable error bar values
    ncs_err = [abs(x-y) for (x,y) in zip(x_ncs,x2)]
    no_ncs_err = [abs(x-y) for (x,y) in zip(y_no_ncs,y2)]

    # Convert lists to numpy arrays
    x_ncs = np.array(x_ncs)
    y_no_ncs = np.array(y_no_ncs)

    # First illustrate basic pyplot interface, using defaults where possible.
    plt.figure()
    colors = []
    # delta_errs = [(x-y) for (x,y) in zip (x_ncs,y_no_ncs)]
    delta_errs = [(x-y) for (x,y) in zip (ncs_err,no_ncs_err)]

    point_size = lambda x: 5 + 1000 * abs(x)
    s = [point_size(x) for x in delta_errs]
    for i,delta_err in enumerate(delta_errs):
      if x > 0:
        colors.append('y')
        print 'Smaller (r_work - r_free) without NCS: ',pdb_code[i]
      else:
        colors.append('b')

    # add points for size reference : Difference of 0, 0.05, 0.1
    s.extend([point_size(0.1),point_size(0.05),point_size(0)])
    colors.extend(['g','g','g'])
    x_ncs = np.append(x_ncs,[.02,.02,.02])
    d = maxval - 0.55
    y_no_ncs = np.append(y_no_ncs,[.4 + d,.45 + d,.5 + d])
    plt.scatter(x_ncs,y_no_ncs,s,c=colors)
    plt.plot([0,maxval],[0,maxval])
    plt.xlabel('R_work Refinement, strict-ncs and transforms',fontsize=14)
    plt.ylabel('R_work Refinement, without strict-ncs',fontsize=14)
    plt.title('Transform and strict-NCS influence on (R_work - R_free)',fontsize=14)
    # plot reference point, to indicate the meaning of size
    plt.text(0.035,0.49 + d, 'The same r_work - r_free value',fontsize=14)
    plt.text(0.035,0.44 + d, '0.05 difference',fontsize=14)
    plt.text(0.035,0.39 + d, '0.1 difference',fontsize=14)

    plt.xlim(0,maxval)
    plt.ylim(0,maxval)
    plt.savefig('transform_refinement.png',transparent=False)
    plt.show()

  def plot_results_4(self):
    """
    Plot a result summery:
    r-free (refined using strict-ncs and Transforms) vs.
    r-free (refined without strict-ncs)
    """
    s = self.sort_inp
    # Collect data points that have results with and without strict_ncs
    x_ncs = []; y_no_ncs =[]
    pdb_code = []
    print len(self.data_records_transform_dict)
    for x in self.data_records_transform_dict:
      if self.data_records_without_strict_ncs_dict.has_key(x):
        x_ncs.append(
          self.data_records_transform_dict[x][s.r_free_asu_final])
        y_no_ncs.append(
          self.data_records_without_strict_ncs_dict[x][s.r_free_asu_final])
        pdb_code.append(x)

    # get larges r-values for the 45 degrees line
    maxval = max(x_ncs + y_no_ncs) *1.05

    # Convert lists to numpy arrays
    x_ncs = np.array(x_ncs)
    y_no_ncs = np.array(y_no_ncs)

    # First illustrate basic pyplot interface, using defaults where possible.
    plt.figure()

    plt.plot([0,maxval],[0,maxval])
    plt.plot(x_ncs,y_no_ncs,'o')
    plt.xlabel('With strict-ncs and transforms',fontsize=14)
    plt.ylabel('Without strict-ncs',fontsize=14)
    plt.title('R-Free after refinement',fontsize=14)

    plt.xlim(0,maxval)
    plt.ylim(0,maxval)
    plt.savefig('r_free.png',transparent=False)
    plt.show()

  def plot_results_5(self):
    """
    Plot a result summery:
    r_free - r_work (refined using strict-ncs and Transforms) vs.
    r_free - r_work (refined without strict-ncs)
    """
    s = self.sort_inp
    # Collect data points that have results with and without strict_ncs
    x_ncs = []; y_no_ncs =[]
    x2 = []; y2 = []
    pdb_code = []
    print len(self.data_records_transform_dict)
    for x in self.data_records_transform_dict:
      if self.data_records_without_strict_ncs_dict.has_key(x):
        x_ncs.append(self.data_records_strict_ncs_dict[x]
                           [s.r_work_asu_final])
        y_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                                  [s.r_work_asu_final])
        x2.append(self.data_records_strict_ncs_dict[x]
                      [s.r_free_asu_final])
        y2.append(self.data_records_without_strict_ncs_dict[x]
                      [s.r_free_asu_final])
        pdb_code.append(x)

    # variable error bar values
    ncs_err = [abs(x-y) for (x,y) in zip(x_ncs,x2)]
    no_ncs_err = [abs(x-y) for (x,y) in zip(y_no_ncs,y2)]

    # Convert lists to numpy arrays
    x_ncs = np.array(ncs_err)
    y_no_ncs = np.array(no_ncs_err)

    # get larges r-values for the 45 degrees line
    maxval = max(x_ncs + y_no_ncs) *1.05
    maxval = 0.2

    # First illustrate basic pyplot interface, using defaults where possible.
    plt.figure()

    plt.plot([0,maxval],[0,maxval])
    plt.plot(x_ncs,y_no_ncs,'o')
    plt.xlabel('With strict-ncs and transforms',fontsize=14)
    plt.ylabel('Without strict-ncs',fontsize=14)
    plt.title('(R_Free - R_work) after refinement',fontsize=14)

    plt.xlim(0,maxval)
    plt.ylim(0,maxval)
    plt.savefig('r_free-r_work.png',transparent=False)
    plt.show()

  def plot_results_6(self):
    """
    Plot a result summery:
    r-free (refined using strict-ncs) vs.
    r-free (refined without strict-ncs)

    ncs_err , no_ncs_err are the difference between r-work and r-free for the two
    refinement methods

    The size of the circles indicate the size of (ncs_err - no_ncs_err). It is blue
    when the strict_ncs refinement make the (r_work - r_free) smaller.
    """
    rec_num = self.sort_inp
    # Collect data points that have results with and without strict_ncs
    x_ncs = []; y_no_ncs =[]
    x2 = []; y2 = []
    time_ncs = []; time_no_ncs = []
    pdb_code = []
    for x in self.data_records_strict_ncs_dict:
      if self.data_records_without_strict_ncs_dict.has_key(x):
        x_ncs.append(self.data_records_strict_ncs_dict[x]
                           [rec_num.r_work_asu_final])
        y_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                                  [rec_num.r_work_asu_final])
        x2.append(self.data_records_strict_ncs_dict[x]
                      [rec_num.r_free_asu_final])
        y2.append(self.data_records_without_strict_ncs_dict[x]
                      [rec_num.r_free_asu_final])
        time_ncs.append(self.data_records_strict_ncs_dict[x]
                      [rec_num.time])
        time_no_ncs.append(self.data_records_without_strict_ncs_dict[x]
                      [rec_num.time])
        pdb_code.append(x)

    # variable error bar values
    ncs_err = [abs(x-y) for (x,y) in zip(x_ncs,x2)]
    no_ncs_err = [abs(x-y) for (x,y) in zip(y_no_ncs,y2)]

    # get larges r-values for the 45 degrees line
    maxval = max(x2 + y2) *1.05

    # Convert lists to numpy arrays
    x2 = np.array(x2)
    y2 = np.array(y2)
    ncs_err = np.array(ncs_err)
    no_ncs_err = np.array(no_ncs_err)

    # First illustrate basic pyplot interface, using defaults where possible.
    plt.figure()
    colors = []
    # delta_errs = [(x-y) for (x,y) in zip (x_ncs,y_no_ncs)]
    delta_errs = [(x-y) for (x,y) in zip (ncs_err,no_ncs_err)]
    point_size = lambda x: 5 + 1000 * abs(x)
    s = [point_size(x) for x in delta_errs]
    for i,delta_err in enumerate(delta_errs):
      if delta_err > 0:
        colors.append('y')
        print 'Smaller (r_work - r_free) without NCS: ',pdb_code[i]
      else:
        colors.append('b')

    # add points for size reference : Difference of 0, 0.05, 0.1
    s.extend([point_size(0.1),point_size(0.05),point_size(0)])
    colors.extend(['g','g','g'])
    x2 = np.append(x2,[.02,.02,.02])
    d = maxval - 0.55
    p1_y = 0.48
    dy = 0.03
    dr = 0.01
    y2 = np.append(y2,[p1_y - 2*dy + d,p1_y - dy + d,p1_y + d])
    plt.scatter(x2,y2,s,c=colors)
    plt.plot([0,maxval],[0,maxval])
    plt.xlabel('R-free with strict-ncs',fontsize=14)
    plt.ylabel('R-free without strict-ncs',fontsize=14)
    plt.title('strict-ncs influence on (R_work - R_free)',fontsize=14)
    # plot reference point, to indicate the meaning of size
    plt.text(0.010,p1_y + dy + d, '(R_free - R_work) difference',fontsize=16)
    plt.text(0.035,p1_y + d - dr, '0.00',fontsize=14)
    plt.text(0.035,p1_y - dy + d - dr, '0.05 ',fontsize=14)
    plt.text(0.035,p1_y - 2*dy + d - dr, '0.10 ',fontsize=14)

    plt.xlim(0,maxval)
    plt.ylim(0,maxval)
    plt.savefig('r_free_bubble_plot.png',transparent=False)
    plt.show()

  def save_csv_table_to_file(self):
    """
    Save a the collected data in a csv file using the record order as
    specified in self.cols_names_for_table1
    """
    print 'Running ',sys._getframe().f_code.co_name
    file_name = 'ncs_refinement_results_v2.csv'
    f = open(file_name,'w')
    f.write(','.join(self.cols_names_for_table1))
    f.write('\n')
    for rec in self.data_records:
      rec = [str(x) for x in rec]
      rec += '\n'
      f.write(','.join(rec))
    f.close()
    print 'data was saved to: ',file_name

  def save_csv_table2_to_file(self):
    """
    Save a the collected data in a csv file:

    The data in the file is one line per test with the records order as
    specified in self.cols_names_for_table2
    """
    print 'Running ',sys._getframe().f_code.co_name
    file_name = 'ncs_refinement_results_table2_v2.csv'
    table_title = self.cols_names_for_table2
    f = open(file_name,'w')
    f.write(','.join(self.cols_names_for_table2))
    f.write('\n')
    for rec in self.data_records_long:
      rec = [str(x) for x in rec]
      rec += '\n'
      f.write(','.join(rec))
    f.close()
    print 'data was saved to: ',file_name

  def save_csv_selected_files(self,pdb_file_list):
    """
    Create a CSV of selected files in a modified format
    """
    print 'Running ',sys._getframe().f_code.co_name
    ur = self.sort_out
    file_name = 'selected_files_v2.csv'
    records_to_use = self.cols_names_for_table2[:ur.r_free_asu_transform]
    table_title = ['PDB','n','Year','Res.','Comp','Solvent',
                   'D/P ASU','D/P NCS',
                   'Header','Initial','W/O NCS','W NCS','Transform']

    f = open(file_name,'w')
    f.write(','.join(table_title))
    f.write('\n')
    for rec in self.data_records_long:
      if rec[ur.pdb_code] in pdb_file_list:
        rec = [str(x) for x in rec]
        outstr = rec[0:8]
        outstr += ['/'.join(rec[ur.r_work_pdb_header:ur.r_free_pdb_header+1]),
                  '/'.join(rec[ur.r_work_asu_init:ur.r_free_asu_init+1]),
                  '/'.join(rec[ur.r_work_asu_no_ncs:ur.r_free_asu_no_ncs+1]),
                  '/'.join(rec[ur.r_work_asu_ncs:ur.r_free_asu_ncs+1]),
                  '/'.join(rec[ur.r_work_asu_transform:ur.r_free_asu_transform+1])]
        outstr += '\n'
        f.write(','.join(outstr))
    f.close()
    print 'data was saved to: ',file_name

  def get_list_of_unprocessed_files(self):
    """
    Check which of the PDB files, from our initial list, are not included in
    the results
    """
    print 'Running ',sys._getframe().f_code.co_name
    pdb_code_set = {
      '3dar', '1vcr', '1r2j', '1a37', '1llc', '1tnv', '1tdi', '1w39', '1ny7',
      '1ddl', '1c8n', '2bfu', '4gmp', '3vbr', '3vbu', '3vbo', '4jgy', '3es5',
      '3nop', '3not', '3nou', '3bcc', '1bcc', '1z7s', '6msf', '2iz8', '7msf',
      '2izn', '2c50', '2c51', '2iz9', '2c4y', '2c4z', '5msf', '2c4q', '2bu1',
      '3raa', '3oah', '3ra2', '3ra9', '3ra8', '3ra4', '3qpr', '1ei7', '1a34',
      '3chx', '2wbh', '2fz1', '2fz2', '2gh8', '1wcd', '3fbm', '4gb3', '1laj',
      '3vbh', '1dzl', '3hag', '4iv3', '1js9', '3n7x', '4gh4', '4jgz', '3tn9',
      '4iv1', '1vb2', '1vb4', '1vak', '3s4g', '2buk', '1x36', '4bcu', '1b35',
      '2wzr', '1k5m', '2bq5', '1zba', '1pgw', '3vbs', '1x35', '3vbf', '1pgl',
      '4fsj', '4fte', '4fts', '2e0z', '4ftb', '2w4y', '2w4z', '2qzv', '3vdd',
      '3p0s', '1qjx', '1qjy', '1qju', '3r0r', '2bs1', '2ztn', '1x9t', '2zzq',
      '1x9p', '4aqq', '1za7', '4ar2', '2wws', '2xpj', '4hl8', '3ntt', '2vf1',
      '3ux1', '2xgk', '2izw', '3cji', '4gbt', '2vq0', '4g93', '2g34', '2qij',
      '2g33', '1f2n', '4g0r', '1ng0', '2ws9', '2xbo', '2wff', '1wce', '1dwn',
      '2vf9', '3zfe', '3zff', '3zfg', '2x5i', '1h8t', '3lob', '4ang', '2gtl',
      '2qqp', '1f8v', '1m1c', '1lp3', '4aed', '3e8k', '1uf2', '1ohg', '1ohf',
      '3s6p', '3kz4', '4f5x', '1vsz'}

    without_strict_ncs_set=set(self.data_records_without_strict_ncs_dict.keys())
    with_strict_ncs_set=set(self.data_records_strict_ncs_dict.keys())
    with_transform_set=set(self.data_records_transform_dict.keys())
    print '\n  files with problems'
    print '-----------------------------------'
    outdata1 = pdb_code_set - without_strict_ncs_set
    outdata2 = pdb_code_set - with_strict_ncs_set
    outdata3 = pdb_code_set - with_transform_set
    outdata = outdata3.union(outdata2)
    outdata = list(outdata.union(outdata1))

    n = 5
    l = len(outdata)
    print 'Total number of files: ',len(pdb_code_set)
    print(list(outdata))
    for i in range(0,l,n):
      e = min(i+n,l)
      s = ['{}']*(e-i)
      s = ', '.join(s)
      print s.format(*outdata[i:e])

    print '\nGood files: ', len(pdb_code_set) - len(outdata)

  def filter_best_1(self):
    """
    Filter files with best (r_work_asu_init - r_work_asu_transform)
    """
    print '\nFilter files with best (r_work_asu_init - r_work_asu_transform)'
    print '---------------------------------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x2,x1,use_data = self.final_vs_init(rec)
      if use_data:
        if rec[use_rec.r_work_asu_transform] != \
                rec[use_rec.r_free_asu_transform]:
          temp.extend([[round(x1-x2,4),rec]])
    temp.sort()
    for rec in temp[-15:]:
      print rec
      file_list.append(rec[1][0])
    print set([x[1][0] for x in temp])

  def filter_best_2(self):
    """
    Filter files with r_free_asu_transform == r_work_asu_transform
    """
    print '\nFilter files with r_free_asu_transform == r_work_asu_transform'
    print '--------------------------------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x1 =  rec[use_rec.r_free_asu_transform]
      y1 =  rec[use_rec.r_work_asu_transform]
      if  isinstance(x1,float):
        if x1 == y1:
          temp.extend([[x1,rec]])
    temp.sort()
    for rec in temp[-10:]:
      file_list.append(rec[1][0])
      print rec
    print set([x[1][0] for x in temp])

  def filter_best_3(self):
    """
    Filter files with best delta(r_work - r_free) improvement
    """
    print '\nFilter files with best delta(r_work - r_free) improvement'
    print '---------------------------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x1 =  rec[use_rec.r_work_asu_ncs]
      x2 =  rec[use_rec.r_free_asu_ncs]
      y1 =  rec[use_rec.r_work_asu_transform]
      y2 =  rec[use_rec.r_free_asu_transform]
      if  isinstance(x1,float) and isinstance(y1,float):
        temp.extend([[round(abs(x1-x2) - abs(y1-y2),4),rec]])
    temp.sort()
    for rec in temp[-10:]:
      file_list.append(rec[1][0])
      print rec
    print set([x[1][0] for x in temp])

  def filter_best_4(self):
    """
    Collect files with final R-work larger than initial
    """
    print '\nFilter files with final R-work larger than initial'
    print '--------------------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x1 =  rec[use_rec.r_work_asu_init]
      y1 =  rec[use_rec.r_work_asu_transform]
      if  isinstance(x1,float) and isinstance(y1,float):
        if y1 > x1:
          temp.extend([[round(y1-x1,4),rec]])
    temp.sort()
    for rec in temp[-10:]:
      file_list.append(rec[1][0])
      print rec
    print set([x[1][0] for x in temp])
    print 'number of files: ',len(temp)

  def r_free_better_without_ncs(self):
    """
    Collect/show structures where the r_free_better_without_ncs

    Argument:
    rec : list of all refinement data
    Return:
    f,i : (float,float) R-free: final,initial
    use_data : (bool) When True, both values are floats
    """
    print '\nr_free_better_without_ncs'
    print '--------------------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x1 =  rec[use_rec.r_free_asu_no_ncs]
      y1 =  rec[use_rec.r_free_asu_transform]
      if  isinstance(x1,float) and isinstance(y1,float):
        if y1 > x1:
          temp.extend([[round(y1-x1,4),rec]])
    temp.sort()
    for rec in temp[-10:]:
      file_list.append(rec[1][0])
      print rec
    print set([x[1][0] for x in temp])
    print 'number of files: ',len(temp)

  def big_change_in_transforms(self,print_all=False, print_n=5):
    """
    Collect files with big change in transformations
    """
    print '\nCollect files with big change in transformations'
    print '------------------------------------------------'
    list_of_files = []
    set_of_files = set()
    for k,v in self.transform_records_dict.iteritems():
      for rec in v:
        m = max(rec)
        if m > 0.2:
          list_of_files.append([m,self.data_records_dict[k]])
          set_of_files.add(k)
    list_of_files.sort()
    list_of_files.reverse()
    printed_files = []
    print 'number of files with large change in transforms: ',len(list_of_files)
    if print_all: print_n = len(list_of_files)
    i = 0
    for l in list_of_files:
      if i == print_n: break
      if l[1][0] not in printed_files:
        print l
        i += 1
        printed_files.append(l[1][0])
    print printed_files[:print_n]

  def filter_best_6(self):
    """
    Collect files with R_work (initial - Header) > 0.2
    """
    print '\nCollect files with R_work (initial - Header) > 0.2'
    print '--------------------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x1 =  rec[use_rec.r_work_asu_init]
      y1 =  rec[use_rec.r_work_pdb_header]
      if  isinstance(x1,float) and isinstance(y1,float):
        if (y1 != -1) and (x1 - y1) > 0.2:
          temp.extend([[round(abs(y1-x1),4),rec]])
    temp.sort()
    for rec in temp:
      file_list.append(rec[1][0])
      print rec
    print set([x[1][0] for x in temp])
    print 'number of files: ',len(temp)

  def filter_best_7(self):
    """
    Collect files with R_work final > 0.5
    """
    print '\nCollect files with R_work final > 0.5'
    print '-------------------------------------'
    temp = []
    file_list = []
    use_rec = self.sort_out
    for rec in self.data_records_long:
      x1 =  rec[use_rec.r_work_asu_transform]
      if  isinstance(x1,float):
        if x1 > 0.5:
          temp.extend([[round(x1,4),rec]])
    temp.sort()
    for rec in temp:
      file_list.append(rec[1][0])
      print rec
    print set([x[1][0] for x in temp])
    print 'number of files: ',len(temp)

  def get_stat(self,exclude_list=[]):
    """
    Get mean and standard deviation of:
    1) improvement in R-work, All refinements methods compare to initial value
    2) improvement in delta(R_work - R_free)  All refinements methods compare to without NCS
    3) improvement in R-work, All refinements methods compare to without transform

    exclude_list: (list of str) contains list of outliers not to include
    """
    use_rec = self.sort_out
    st1 ='R-work improvement, All refinements methods compare to initial value'
    sf1 = self.final_vs_init
    st2 ='delta(R_work - R_free) improvement, All refinements methods compare to without NCS'
    sf2 = self.delta_work_free
    st3 ='R-work improvement, applying transform refinement'
    sf3 = self.final_vs_ncs
    st4 = 'delta(R_work - R_free) All refinements methods compare to PDB'
    sf4 = self.delta_work_free_pdb
    get_stat_using = [(st1,sf1),(st2,sf2),(st3,sf3),(st4,sf4)]
    for test in get_stat_using:
      i = []
      f = []
      print test[0]
      for rec in self.data_records_long:
        f1,i1,use_data =test[1](rec)
        if use_data:
          i.append(i1)
          f.append(f1)
      i = np.array(i)
      f = np.array(f)
      d = i - f
      print 'Mean : {0:.4f}  STD : {1:.4f}   Number of files {2}'\
        .format(np.mean(d),np.std(d),len(d))

  def final_vs_init(self,rec):
    """
    Compare R-work, All refinements methods compare to initial value

    Argument:
    rec : list of all refinement data
    Return:
    f,i : (float,float) R-work : final,initial
    use_data : (bool) When True, both values are floats
    """
    use_rec = self.sort_out
    f =  rec[use_rec.r_work_asu_transform]
    i =  rec[use_rec.r_work_asu_init]
    use_data = isinstance(f,float) and isinstance(i,float)
    return f,i,use_data

  def final_vs_pdb(self,rec):
    """
    Compare R-work, All refinements methods compare to pdb published value

    Argument:
    rec : list of all refinement data
    Return:
    f,i : (float,float) R-work : final,initial
    use_data : (bool) When True, both values are floats
    """
    use_rec = self.sort_out
    f =  rec[use_rec.r_work_asu_transform]
    i =  rec[use_rec.r_work_pdb_header]
    use_data = isinstance(f,float) and isinstance(i,float)
    return f,i,use_data

  def delta_work_free(self,rec):
    """
    Compare delta(R_work - R_free)  All refinements methods compare to without NCS

    Argument:
    rec : list of all refinement data
    Return:
    f,i : (float,float) delta(R_work - R_free) : final,initial
    use_data : (bool) When True, both values are floats
    """
    use_rec = self.sort_out
    i1 =  rec[use_rec.r_work_asu_no_ncs]
    i2 =  rec[use_rec.r_free_asu_no_ncs]
    f1 =  rec[use_rec.r_work_asu_transform]
    f2 =  rec[use_rec.r_free_asu_transform]
    use_data = isinstance(f1,float) and isinstance(i1,float)
    if use_data:
      f =  abs(f2-f1)
      i =  abs(i2-i1)
    else:
      f = i = None
    return f,i,use_data

  def delta_work_free_pdb(self,rec):
    """
    Compare delta(R_work - R_free)  All refinements methods compare to PDB

    Argument:
    rec : list of all refinement data
    Return:
    f,i : (float,float) delta(R_work - R_free) : final,initial
    use_data : (bool) When True, both values are floats
    """
    use_rec = self.sort_out
    i1 =  rec[use_rec.r_work_pdb_header]
    i2 =  rec[use_rec.r_free_pdb_header]
    f1 =  rec[use_rec.r_work_asu_transform]
    f2 =  rec[use_rec.r_free_asu_transform]
    test1 = isinstance(i1,float) and isinstance(i2,float)
    test2 = test1 and (-1 not in [i1,i2])
    use_data = isinstance(f1,float) and test2
    if use_data:
      f =  abs(f2-f1)
      i =  abs(i2-i1)
    else:
      f = i = None
    return f,i,use_data

  def final_vs_ncs(self,rec):
    """
    Compare R-work, All refinements methods compare to without transform

    Argument:
    rec : list of all refinement data
    Return:
    f,i : (float,float) R-work : final,initial
    use_data : (bool) When True, both values are floats
    """
    use_rec = self.sort_out
    f =  rec[use_rec.r_work_asu_transform]
    i =  rec[use_rec.r_work_asu_ncs]
    use_data = isinstance(f,float) and isinstance(i,float)
    return f,i,use_data

  def filter_by_year(self,rec_list,year):
    """d
    filter a list of [[float,list],[float,list]...] by the year value in list
    Return rec_list with records from 'year' and on
    """
    use_rec = self.sort_out
    return [x for x in rec_list if x[1][use_rec.year] >= year]

  def process_data_according_to_data_type(self,ln):
    """ Process line of data from data files  """
    data = [x.strip() for x in ln.split('|')]
    # convert data values to proper type
    for indx in self.float_type_records:
      data[indx] = float(data[indx])
    for indx in self.int_type_records:
      if data[indx] == 'None': data[indx] = 0
      else: data[indx] = int(data[indx])
    for indx in self.bool_type_records:
      data[indx] = (data[indx] == 'True')
    return data

  def combine_records(self,data):
    """
    Every pdb record can be refined with different options. Collect and
    combine all records related to a pdb record to a single data line
    """
    use_rec = self.sort_inp
    if self.data_records_dict.has_key(data[0]):
      data_record = self.data_records_dict[data[0]]
    else:
      data_record = len(self.cols_names)*[None,]
    if data[use_rec.use_ncs_constraints] and not data[use_rec.use_transforms]:
      for map_from,map_to in self.map_to_ncs.iteritems():
        if not  data_record[map_to]:
          data_record[map_to] = data[map_from]
    if not data[use_rec.use_ncs_constraints] and not data[use_rec.use_transforms]:
      for map_from,map_to in self.map_to_no_ncs.iteritems():
        if not  data_record[map_to]:
          data_record[map_to] = data[map_from]
    if data[use_rec.use_ncs_constraints] and data[use_rec.use_transforms]:
      for map_from,map_to in self.map_to_transform.iteritems():
        if not  data_record[map_to]:
          data_record[map_to] = data[map_from]
    return data_record

def print_warning(ln,pdb_code):
  if ln.startswith('Warning'):
    print '{0}: {1}'.format(pdb_code,ln.strip())
  if ln.startswith('Using pdb file from local machine'):
    print 'Used {} from local machine'.format(pdb_code)
  if ln.startswith('Traceback'):
    print '{0}: {1}'.format(pdb_code,'Problem processing')

if __name__=='__main__':
  # fn = r'C:\Phenix\Dev\Work\work\NCS\junk\pdb_test\Refinement_data_9_12_2014'
  # fn = "/net/cci/youval/Work/work/NCS/junk/pdb_test/queue_job"
  fn = r"C:\Phenix\Dev\Work\work\NCS\junk\pdb_test\queue_job"

  path_to_log_files = fn
  current_path = os.getcwd()
  os.chdir(path_to_log_files)
  process_results = results_collection()

  # need to read_filenames only for new runs, it overwrites existing data files
  process_results.read_filenames()

  # process existing files
  process_results.get_data_from_files()
  process_results.save_csv_table_to_file()
  process_results.save_csv_table2_to_file()
  print '-'*30

  # plot
  process_results.plot_results_1()
  process_results.plot_results_2()
  process_results.plot_results_3()
  process_results.plot_results_4()
  process_results.plot_results_5()
  process_results.plot_results_6()

  # check which files are not processed yet
  process_results.get_list_of_unprocessed_files()

  s1 = ['3raa','2bfu','3qpr','3ntt','1dzl','3s4g','1ei7',
        '2qij','2ztn','3nou','3hag','2g34']
  s2 = ['3raa','2bfu','3qpr','3ntt','1dzl','3s4g','1ei7',
        '2qij','3nou','3hag','2g34',
        '1dwn','1qjy', '1qju', '1a37', '2buk','1vcr','1tnv',
        '3r0r','1vb2','4gbt','2wws','2wff']

  process_results.filter_best_1()
  process_results.filter_best_2()
  process_results.filter_best_3()
  process_results.filter_best_4()
  process_results.big_change_in_transforms()
  process_results.r_free_better_without_ncs()
  process_results.filter_best_6()
  process_results.filter_best_7()
  process_results.save_csv_selected_files(s2)
  process_results.get_stat(exclude_list=[])

  os.chdir(current_path)
  print 'Done...'

max_stars_repo_name: youdar/work

max_issues_repo_name: youdar/work

max_forks_repo_name: youdar/work


text:
# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPT2Tokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ..generation.test_generation_flax_utils import FlaxGenerationTesterMixin
from ..test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask

if is_flax_available():
    import jax
    import jax.numpy as jnp
    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
    import torch

class FlaxGPTNeoModelTester:
    def __init__(
        self,
        parent,
        batch_size=14,
        seq_length=7,
        is_training=True,
        use_input_mask=True,
        use_token_type_ids=False,
        use_labels=True,
        vocab_size=99,
        hidden_size=32,
        num_hidden_layers=4,
        num_attention_heads=4,
        attention_types=[[["global", "local"], 2]],
        intermediate_size=37,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        max_position_embeddings=512,
        window_size=7,
        initializer_range=0.02,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.seq_length = seq_length
        self.is_training = is_training
        self.use_input_mask = use_input_mask
        self.use_token_type_ids = use_token_type_ids
        self.use_labels = use_labels
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.attention_types = attention_types
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.window_size = window_size
        self.initializer_range = initializer_range
        self.scope = None
        self.bos_token_id = vocab_size - 1
        self.eos_token_id = vocab_size - 1
        self.pad_token_id = vocab_size - 1

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)

        input_mask = None
        if self.use_input_mask:
            input_mask = random_attention_mask([self.batch_size, self.seq_length])

        config = GPTNeoConfig(
            vocab_size=self.vocab_size,
            hidden_size=self.hidden_size,
            num_layers=self.num_hidden_layers,
            num_heads=self.num_attention_heads,
            max_position_embeddings=self.max_position_embeddings,
            use_cache=False,
            bos_token_id=self.bos_token_id,
            eos_token_id=self.eos_token_id,
            pad_token_id=self.pad_token_id,
            window_size=self.window_size,
            attention_types=self.attention_types,
        )

        return (config, input_ids, input_mask)

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, input_ids, attention_mask = config_and_inputs
        inputs_dict = {"input_ids": input_ids, "attention_mask": attention_mask}
        return config, inputs_dict

    def check_use_cache_forward(self, model_class_name, config, input_ids, attention_mask):
        max_decoder_length = 20
        model = model_class_name(config)

        past_key_values = model.init_cache(input_ids.shape[0], max_decoder_length)
        attention_mask = jnp.ones((input_ids.shape[0], max_decoder_length), dtype="i4")

        position_ids = jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1)[None, :], (input_ids.shape[0], input_ids.shape[-1] - 1)
        )
        outputs_cache = model(
            input_ids[:, :-1],
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )

        position_ids = jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]], dtype="i4")
        outputs_cache_next = model(
            input_ids[:, -1:],
            attention_mask=attention_mask,
            past_key_values=outputs_cache.past_key_values,
            position_ids=position_ids,
        )

        outputs = model(input_ids)

        diff = np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
        self.parent.assertTrue(diff < 1e-3, msg=f"Max diff is {diff}")

    def check_use_cache_forward_with_attn_mask(self, model_class_name, config, input_ids, attention_mask):
        max_decoder_length = 20
        model = model_class_name(config)

        attention_mask_cache = jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1]))],
            axis=-1,
        )

        past_key_values = model.init_cache(input_ids.shape[0], max_decoder_length)
        position_ids = jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1)[None, :], (input_ids.shape[0], input_ids.shape[-1] - 1)
        )

        outputs_cache = model(
            input_ids[:, :-1],
            attention_mask=attention_mask_cache,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )
        position_ids = jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]], dtype="i4")
        outputs_cache_next = model(
            input_ids[:, -1:],
            past_key_values=outputs_cache.past_key_values,
            attention_mask=attention_mask_cache,
            position_ids=position_ids,
        )

        outputs = model(input_ids, attention_mask=attention_mask)

        diff = np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
        self.parent.assertTrue(diff < 1e-3, msg=f"Max diff is {diff}")

@require_flax
class FlaxGPTNeoModelTest(FlaxModelTesterMixin, FlaxGenerationTesterMixin, unittest.TestCase):

    all_model_classes = (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
    all_generative_model_classes = (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()

    def setUp(self):
        self.model_tester = FlaxGPTNeoModelTester(self)

    def test_use_cache_forward(self):
        for model_class_name in self.all_model_classes:
            config, input_ids, attention_mask = self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward(model_class_name, config, input_ids, attention_mask)

    def test_use_cache_forward_with_attn_mask(self):
        for model_class_name in self.all_model_classes:
            config, input_ids, attention_mask = self.model_tester.prepare_config_and_inputs()
            self.model_tester.check_use_cache_forward_with_attn_mask(
                model_class_name, config, input_ids, attention_mask
            )

    @slow
    def test_batch_generation(self):
        tokenizer = GPT2Tokenizer.from_pretrained("gpt2", pad_token="<|endoftext|>", padding_side="left")
        inputs = tokenizer(["Hello this is a long string", "Hey"], return_tensors="np", padding=True, truncation=True)

        model = FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
        model.do_sample = False
        model.config.pad_token_id = model.config.eos_token_id

        jit_generate = jax.jit(model.generate)

        output_sequences = jit_generate(
            inputs["input_ids"], attention_mask=inputs["attention_mask"], pad_token_id=tokenizer.pad_token_id
        ).sequences

        output_string = tokenizer.batch_decode(output_sequences, skip_special_tokens=True)

        expected_string = [
            "Hello this is a long string of text.\n\nI'm trying to get the text of the",
            "Hey, I'm a little late to the party. I'm going to",
        ]

        self.assertListEqual(output_string, expected_string)

    # overwrite from common since `attention_mask` in combination
    # with `causal_mask` behaves slighly differently
    @is_pt_flax_cross_test
    def test_equivalence_pt_to_flax(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__):
                # prepare inputs
                prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
                pt_inputs = {k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

                # load corresponding PyTorch class
                pt_model_class_name = model_class.__name__[4:]  # Skip the "Flax" at the beginning
                pt_model_class = getattr(transformers, pt_model_class_name)

                batch_size, seq_length = pt_inputs["input_ids"].shape
                rnd_start_indices = np.random.randint(0, seq_length - 1, size=(batch_size,))
                for batch_idx, start_index in enumerate(rnd_start_indices):
                    pt_inputs["attention_mask"][batch_idx, :start_index] = 0
                    pt_inputs["attention_mask"][batch_idx, start_index:] = 1
                    prepared_inputs_dict["attention_mask"][batch_idx, :start_index] = 0
                    prepared_inputs_dict["attention_mask"][batch_idx, start_index:] = 1
                pt_model = pt_model_class(config).eval()
                fx_model = model_class(config, dtype=jnp.float32)

                fx_state = convert_pytorch_state_dict_to_flax(pt_model.state_dict(), fx_model)
                fx_model.params = fx_state

                with torch.no_grad():
                    pt_outputs = pt_model(**pt_inputs).to_tuple()

                fx_outputs = fx_model(**prepared_inputs_dict).to_tuple()
                self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch")
                for fx_output, pt_output in zip(fx_outputs, pt_outputs):
                    self.assert_almost_equals(fx_output[:, -1], pt_output[:, -1].numpy(), 4e-2)

                with tempfile.TemporaryDirectory() as tmpdirname:
                    pt_model.save_pretrained(tmpdirname)
                    fx_model_loaded = model_class.from_pretrained(tmpdirname, from_pt=True)

                fx_outputs_loaded = fx_model_loaded(**prepared_inputs_dict).to_tuple()
                self.assertEqual(
                    len(fx_outputs_loaded), len(pt_outputs), "Output lengths differ between Flax and PyTorch"
                )
                for fx_output_loaded, pt_output in zip(fx_outputs_loaded, pt_outputs):
                    self.assert_almost_equals(fx_output_loaded[:, -1], pt_output[:, -1].numpy(), 4e-2)

    # overwrite from common since `attention_mask` in combination
    # with `causal_mask` behaves slighly differently
    @is_pt_flax_cross_test
    def test_equivalence_flax_to_pt(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__):
                # prepare inputs
                prepared_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
                pt_inputs = {k: torch.tensor(v.tolist()) for k, v in prepared_inputs_dict.items()}

                # load corresponding PyTorch class
                pt_model_class_name = model_class.__name__[4:]  # Skip the "Flax" at the beginning
                pt_model_class = getattr(transformers, pt_model_class_name)

                pt_model = pt_model_class(config).eval()
                fx_model = model_class(config, dtype=jnp.float32)

                pt_model = load_flax_weights_in_pytorch_model(pt_model, fx_model.params)
                batch_size, seq_length = pt_inputs["input_ids"].shape
                rnd_start_indices = np.random.randint(0, seq_length - 1, size=(batch_size,))
                for batch_idx, start_index in enumerate(rnd_start_indices):
                    pt_inputs["attention_mask"][batch_idx, :start_index] = 0
                    pt_inputs["attention_mask"][batch_idx, start_index:] = 1
                    prepared_inputs_dict["attention_mask"][batch_idx, :start_index] = 0
                    prepared_inputs_dict["attention_mask"][batch_idx, start_index:] = 1

                # make sure weights are tied in PyTorch
                pt_model.tie_weights()

                with torch.no_grad():
                    pt_outputs = pt_model(**pt_inputs).to_tuple()

                fx_outputs = fx_model(**prepared_inputs_dict).to_tuple()
                self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch")
                for fx_output, pt_output in zip(fx_outputs, pt_outputs):
                    self.assert_almost_equals(fx_output[:, -1], pt_output[:, -1].numpy(), 4e-2)

                with tempfile.TemporaryDirectory() as tmpdirname:
                    fx_model.save_pretrained(tmpdirname)
                    pt_model_loaded = pt_model_class.from_pretrained(tmpdirname, from_flax=True)

                with torch.no_grad():
                    pt_outputs_loaded = pt_model_loaded(**pt_inputs).to_tuple()

                self.assertEqual(
                    len(fx_outputs), len(pt_outputs_loaded), "Output lengths differ between Flax and PyTorch"
                )
                for fx_output, pt_output in zip(fx_outputs, pt_outputs_loaded):
                    self.assert_almost_equals(fx_output[:, -1], pt_output[:, -1].numpy(), 4e-2)

    @slow
    def test_model_from_pretrained(self):
        for model_class_name in self.all_model_classes:
            model = model_class_name.from_pretrained("EleutherAI/gpt-neo-125M")
            outputs = model(np.ones((1, 1)))
            self.assertIsNotNone(outputs)

max_stars_repo_name: dctelus/transformers

reason: import numpy,import jax


text:
# coding: utf-8

# In[3]:

get_ipython().magic(u'matplotlib inline')
from __future__ import print_function
import numpy as np
import mdtraj as md
t = md.load_frame('.\\trajectory-1.xtc', 9999, top='.\\fs-peptide.pdb')
print(t)

# In[7]:

hbonds = md.baker_hubbard(t, periodic=False)
label = lambda hbond : '%s -- %s' % (t.topology.atom(hbond[0]), t.topology.atom(hbond[2]))
for hbond in hbonds:
    print(label(hbond))

hbond_score = len(hbonds)
print('\n\nH-Bond score:', hbond_score)
# calculate heavy atom to atom contacts
# basically a number that looks like Nij/root(ninj)

# In[8]:

da_distances = md.compute_distances(t, hbonds[:, [0,2]], periodic=False)

# In[14]:

import itertools
import matplotlib.pyplot as plt
color = itertools.cycle(['r', 'b', 'gold'])
for i in [2, 3, 4]:
    plt.hist(da_distances[:, i], color=next(color), label=label(hbonds[i]), alpha=0.5)
plt.legend()
plt.ylabel('Freq');
plt.xlabel('Donor-acceptor distance [nm]')

max_stars_repo_name: yngtodd/molecules-deprecated

max_issues_repo_name: yngtodd/molecules-deprecated

max_forks_repo_name: yngtodd/molecules-deprecated


text:
import numpy as np
from skcv import skcv_knn, skcv_rls
from lposcv import lposcv_knn, lposcv_rls

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""

 --- EXAMPLE SCRIPT ---

 DESCRIPTION:
 - In this script we will test and visualize
both SKCV and LPO-SCV methods with suitable data sets.
SKCV refers to the spatial K-fold cross-validation method,
whereas LPO-SCV corresponds to the leave-pair-out spatial
cross-validation method.

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""

if __name__=="__main__":

    ########################################################
    ## STEP 1: SKCV EXPERIMENT
    ##
    ## - In this part we will implement SKCV-analysis for
    ## open natural data. The output variable of the analysis
    ## is soil water permeability.
    ##
    ########################################################
    # First we load data coordinates
    coordinates_skcv =  np.loadtxt("data_sets/coordinates_skcv.txt", delimiter=",")
    # Predictor features
    Xdata_skcv =  np.loadtxt("data_sets/Xdata_skcv.txt", delimiter=",")
    # Water permeability output variable
    Ydata_skcv =  np.loadtxt("data_sets/Ydata_skcv.txt", delimiter=",")
    # Next we set analysis parameters. We try first ridge regression as the prediction model.
    regparam = 1. # Regularization parameter
    dead_zone_radiuses = range(0, 2500, 500) # Produce SKCV for dead zone radiuses 0, 500, ..., 2000 meters.
    number_of_folds = 10 # Do 10-fold cross-validation
    # If visualization is wanted, set 'visualization = True'. Visualization will result in significantly longer calculation time.
    visualization = True
    # Run the analysis
    skcv_rls(coordinates_skcv, Xdata_skcv, Ydata_skcv, number_of_folds, dead_zone_radiuses, regparam, visualization)
    # Next we try the same analysis, but for k-nearest neighbor
    k = 10 # Number of neighbors
    # Run the analysis
    skcv_knn(coordinates_skcv, Xdata_skcv, Ydata_skcv, number_of_folds, dead_zone_radiuses, k, visualization)

    ########################################################
    ## STEP 2: LPO-SCV EXPERIMENT
    ##
    ## - In this part we will implement LPO-SCV-analysis for
    ## natural data sets. The output variable of the analysis
    ## is soil gold occurrence.
    ##
    ########################################################
    # First we load data coordinates
    coordinates_lposcv =  np.loadtxt("data_sets/coordinates_lposcv.txt", delimiter=",")
    # Predictor features
    Xdata_lposcv =  np.loadtxt("data_sets/Xdata_lposcv.txt", delimiter=",")
    Xdata_lposcv = Xdata_lposcv[:,2:]
    # Output variable
    Ydata_lposcv =  np.loadtxt("data_sets/Ydata_lposcv.txt", delimiter=",")
    # Next we set analysis parameters. We try first ridge regression as the prediction model.
    regparam = 1. # Regularization parameter
    dead_zone_radiuses = range(0, 30000, 5000) # Produce SKCV for dead zone radiuses 0, 1000, ..., 4000 meters.
    # If visualization is wanted, set 'visualization = True'. Visualization will result in significantly longer calculation time.
    visualization = False
    # Run the analysis
    lposcv_rls(coordinates_lposcv, Xdata_lposcv, Ydata_lposcv, dead_zone_radiuses, regparam, visualization)
    # Next we try the same analysis, but for k-nearest neighbor
    k = 10 # Number of neighbors
    # Run the analysis
    lposcv_knn(coordinates_lposcv, Xdata_lposcv, Ydata_lposcv, dead_zone_radiuses, k, visualization)

max_stars_repo_name: jjepsuomi/LPO-SCV

max_issues_repo_name: jjepsuomi/LPO-SCV

max_forks_repo_name: jjepsuomi/LPO-SCV


text:
import numpy as np

from velocileptors.LPT.velocity_moments_fftw import VelocityMoments

class FourierStreamingModel(VelocityMoments):
    '''
    Class to calculate the redshift space power spectrum in the Fourier streaming model.

    Inherits the VelocityMoments class which itself inherits the CLEFT class.

    Here we compute up to the first two moments v(k) and sigma(k), in addition
    to a counterterm ansatz for the third moment, which is an excellent approximation to the data.

    The full one-loop expression can be computed in the moment expansion class if desired,
    but we note that the inclusion of higher moments did not noticeably improve agreement
    with data in n-body simulations.

    '''

    def __init__(self, *args, kmin = 1e-3, kmax = 3, nk = 100, **kw):
        '''
         Same keywords and arguments as the other two classes for now.
        '''

        # Setup ffts etc.
        VelocityMoments.__init__(self, *args, **kw)

        self.kmin, self.kmax, self.nk = kmin, kmax, nk
        self.kv = np.logspace(np.log10(kmin), np.log10(kmax), nk); self.nk = nk
        self.weight = self.kv**3 / (2*np.pi**2)

        self.setup_velocity_moments()

    def setup_velocity_moments(self):
        self.make_ptable(kmin = self.kmin, kmax = self.kmax, nk = self.nk)
        self.make_vtable(kmin = self.kmin, kmax = self.kmax, nk = self.nk)
        self.make_spartable(kmin = self.kmin, kmax = self.kmax, nk = self.nk)
        self.make_stracetable(kmin = self.kmin, kmax = self.kmax, nk = self.nk)
        self.convert_sigma_bases(basis='Legendre')

    def compute_cumulants(self, pars):
        '''
        Calculate velocity moments and turn into cumulants.
        The pars format is [b1, b2, bs, b3, alpha, alpha_v, alpha_s0, alpha_s2, sn, sv, sigma0_stoch]

        '''
        # Compute each moment
        b1, b2, bs, b3, alpha, alpha_v, alpha_s0, alpha_s2, sn, sv, sigma0_stoch = pars

        kv, pk = self.combine_bias_terms_pk(b1,b2,bs,b3,alpha,sn)
        kv, vk = self.combine_bias_terms_vk(b1,b2,bs,b3,alpha_v,sv)
        kv, s0, s2 = self.combine_bias_terms_sk(b1,b2,bs,b3,alpha_s0,alpha_s2,sigma0_stoch,basis='Legendre')

        # Turn into cumulants; split into coefficients of LOS angle mu
        self.one_plus_delta = 1 + self.weight * pk

        # first cumulant
        # this is imaginary and goes into the expansion as exp(- (k*mu) im_c1 * mu ...)
        self.im_c1 = self.weight * vk / self.one_plus_delta

        # second cumulant: note that sigma_zz = sigma0 + sigma2 * 3/2 (mu^2 - 1/3) = (sigma0 -1/2 sigma 2) + 3/2 * sigma2 * mu^2
        self.c2_0 = self.weight * (s0 - 0.5*s2)/self.one_plus_delta
        self.c2_2 = self.weight * (1.5*s2)/self.one_plus_delta + self.im_c1**2

    def compute_redshift_space_power_at_mu(self, pars, f, mu, counterterm_c3=0):

        self.compute_cumulants(pars)

        mu2 = mu**2
        expon = -f * self.kv * mu2 * self.im_c1 - 0.5 * f**2 * self.kv**2 * mu2 * (self.c2_0 + self.c2_2*mu2)
        expon += counterterm_c3/6. * self.weight * self.kv**2 * mu**4 * self.pktable[:,-1] / self.one_plus_delta

        return self.kv, (self.one_plus_delta * np.exp(expon) - 1.)/self.weight

    def compute_redshift_space_power_multipoles(self, pars, f, counterterm_c3=0, ngauss=4):

        # Generate the sampling
        mus, ws = np.polynomial.legendre.leggauss(2*ngauss)
        mus_calc = mus[0:ngauss]

        L0 = np.polynomial.legendre.Legendre((1))(mus)
        L2 = np.polynomial.legendre.Legendre((0,0,1))(mus)
        L4 = np.polynomial.legendre.Legendre((0,0,0,0,1))(mus)

        self.pkmutable = np.zeros((len(mus),self.nk))

        for ii, mu in enumerate(mus_calc):
            self.pkmutable[ii,:] = self.compute_redshift_space_power_at_mu(pars,f,mu,counterterm_c3=counterterm_c3)[1]

        self.pkmutable[ngauss:,:] = np.flip(self.pkmutable[0:ngauss],axis=0)

        self.p0ktable = 0.5 * np.sum((ws*L0)[:,None]*self.pkmutable,axis=0)
        self.p2ktable = 2.5 * np.sum((ws*L2)[:,None]*self.pkmutable,axis=0)
        self.p4ktable = 4.5 * np.sum((ws*L4)[:,None]*self.pkmutable,axis=0)

        return self.kv, self.p0ktable, self.p2ktable, self.p4ktable

max_stars_repo_name: kokron/velocileptors

max_issues_repo_name: kokron/velocileptors

max_forks_repo_name: kokron/velocileptors


text:
#!/usr/bin/env python -i
# preceding line should have path for Python on your machine

# simple.py
# Purpose: mimic operation of examples/COUPLE/simple/simple.cpp via Python

# Serial syntax: simple.py in.lammps
#                in.lammps = LAMMPS input script

# Parallel syntax: mpirun -np 4 simple.py in.lammps
#                  in.lammps = LAMMPS input script
# also need to uncomment either Pypar or mpi4py sections below

from __future__ import print_function
import sys
import numpy as np
import ctypes

# parse command line

argv = sys.argv
if len(argv) != 2:
  print("Syntax: simple.py in.lammps")
  sys.exit()

infile = sys.argv[1]

me = 0

# uncomment this if running in parallel via Pypar
#import pypar
#me = pypar.rank()
#nprocs = pypar.size()

# uncomment this if running in parallel via mpi4py
#from mpi4py import MPI
#me = MPI.COMM_WORLD.Get_rank()
#nprocs = MPI.COMM_WORLD.Get_size()

from lammps import lammps
lmp = lammps()

# run infile one line at a time

lines = open(infile,'r').readlines()
for line in lines: lmp.command(line)

# run 10 more steps
# get coords from LAMMPS
# change coords of 1st atom
# put coords back into LAMMPS
# run a single step with changed coords

lmp.command("run 10")
x = lmp.gather_atoms("x",1,3)
v = lmp.gather_atoms("v",1,3)
epsilon = 0.1
x[0] += epsilon
lmp.scatter_atoms("x",1,3,x)
lmp.command("run 1");

# extract force on single atom two different ways

f = lmp.extract_atom("f",3)
print("Force on 1 atom via extract_atom: ",f[0][0])

fx = lmp.extract_variable("fx","all",1)
print("Force on 1 atom via extract_variable:",fx[0])

# use commands_string() and commands_list() to invoke more commands

strtwo = "run 10\nrun 20"
lmp.commands_string(strtwo)

cmds = ["run 10","run 20"]
lmp.commands_list(cmds)

# delete all atoms
# create_atoms() to create new ones with old coords, vels
# initial thermo should be same as step 20

natoms = lmp.get_natoms()
atype = natoms*[1]

lmp.command("delete_atoms group all");
lmp.create_atoms(natoms,None,atype,x,v);
lmp.command("run 10");

############
# test of new gather/scatter and box extract/reset methods
# can try this in parallel and with/without atom_modify sort enabled

lmp.command("write_dump all custom tmp.simple id type x y z fx fy fz");

x = lmp.gather_atoms("x",1,3)
f = lmp.gather_atoms("f",1,3)

if me == 0: print("Gather XF:",x[3],x[9],f[3],f[9])

ids = lmp.gather_atoms_concat("id",0,1)
x = lmp.gather_atoms_concat("x",1,3)
f = lmp.gather_atoms_concat("f",1,3)

if me == 0: print("Gather concat XF:",ids[0],ids[1],x[0],x[3],f[0],f[3])

ids = (2*ctypes.c_int)()
ids[0] = 2
ids[1] = 4
x = lmp.gather_atoms_subset("x",1,3,2,ids)
f = lmp.gather_atoms_subset("f",1,3,2,ids)

if me == 0: print("Gather subset XF:",x[0],x[3],f[0],f[3])

x[0] = -1.0
x[1] = 0.0
x[2] = 0.0
x[3] = -2.0
x[4] = 0.0
x[5] = 0.0
ids[0] = 100
ids[1] = 200
lmp.scatter_atoms_subset("x",1,3,2,ids,x)

x = lmp.gather_atoms("x",1,3)
if me == 0: print("Gather post scatter subset:",
                  x[3],x[9],x[297],x[298],x[299],x[597],x[598],x[599])

boxlo,boxhi,xy,yz,xz,periodicity,box_change = lmp.extract_box()
if me == 0: print("Box info",boxlo,boxhi,xy,yz,xz,periodicity,box_change)

lmp.reset_box([0,0,0],[10,10,8],0,0,0)

boxlo,boxhi,xy,yz,xz,periodicity,box_change = lmp.extract_box()
if me == 0: print("Box info",boxlo,boxhi,xy,yz,xz,periodicity,box_change)

# uncomment if running in parallel via Pypar
#print("Proc %d out of %d procs has" % (me,nprocs), lmp)

# uncomment if running in parallel via mpi4py
#print("Proc %d out of %d procs has" % (me,nprocs), lmp)