text:
# hotel_cleaner.py
# author: Trevor Kinsey
# date: 2021-01-19

'''This script cleans and wrangles the hotels.csv file to be used
   in a visualization dashboard app

Usage: python hotel_cleaner.py

'''
import numpy as np
import pandas as pd

# import the data
hotels = pd.read_csv("data/raw/hotels.csv")

# create new columns from other columns
months = ["January", "February", "March", "April",
          "May", "June", "July", "August", "September",
          "October", "November", "December"]
hotels["arrival_date_month"] = hotels["arrival_date_month"].replace(months,[1,2,3,4,5,6,7,8,9,10,11,12])
hotels["Arrival date"] = pd.to_datetime(hotels.arrival_date_year*10000 + hotels.arrival_date_month*100 + hotels.arrival_date_day_of_month,
                                        format = '%Y%m%d')
hotels["Arrival day of week"] = hotels["Arrival date"].dt.dayofweek
hotels["Arrival day of week"] = hotels["Arrival day of week"].replace([0,1,2,3,4,5,6],["Mon", "Tues", "Wed", "Thur", "Fri", "Sat", "Sun"])
hotels["Total nights"] = hotels["stays_in_weekend_nights"] + hotels["stays_in_week_nights"]
# drop unused columns
hotels = hotels.drop(columns=['agent', 'company', 'lead_time',
                              'market_segment', 'distribution_channel',
                              'is_repeated_guest', 'previous_cancellations',
                              'previous_bookings_not_canceled', 'reserved_room_type',
                              'assigned_room_type','deposit_type',
                              'days_in_waiting_list', 'customer_type', 'reservation_status',
                              'reservation_status_date', 'meal'], )
# Change values to make more readable
hotels["hotel"] = hotels["hotel"].replace(["Resort Hotel", "City Hotel"], ["Resort", "City"])
# change column names to make more readable
hotels.columns = ['Hotel type', 'Cancelled', 'Arrival year',
                  'Arrival month', 'Arrival week', 'Arrival day', 'Weekend nights',
                  'Week nights', 'Adults', 'Children', 'Babies', 'Country of origin',
                  'Booking changes', 'Average daily rate', 'Required parking spaces',
                  'Special requests', "Arrival date", "Arrival day of week", 'Total nights']

# # create date index column for date slider
# hotels["arrival_date_index"] = pd.to_numeric(hotels["Arrival date"])/60/60/24/10**9
# start_date = hotels["arrival_date_index"][0]
# hotels["arrival_date_index"] -= start_date

#  save to file
hotels.to_csv("data/processed/clean_hotels.csv", index = False)


text:
import math
import numpy as np
from sklearn.neighbors import NearestNeighbors

def euclidean_distance(point1, point2):
    """
    Euclidean distance between two points.
    :param point1: the first point as a tuple (a_1, a_2, ..., a_n)
    :param point2: the second point as a tuple (b_1, b_2, ..., b_n)
    :return: the Euclidean distance
    """
    a = np.array(point1)
    b = np.array(point2)

    return np.linalg.norm(a - b, ord=2)

def point_based_matching(point_pairs):
    """
    This function is based on the paper "Robot Pose Estimation in Unknown Environments by Matching 2D Range Scans"
    by F. Lu and E. Milios.

    :param point_pairs: the matched point pairs [((x1, y1), (x1', y1')), ..., ((xi, yi), (xi', yi')), ...]
    :return: the rotation angle and the 2D translation (x, y) to be applied for matching the given pairs of points
    """

    x_mean = 0
    y_mean = 0
    xp_mean = 0
    yp_mean = 0
    n = len(point_pairs)

    if n == 0:
        return None, None, None

    for pair in point_pairs:

        (x, y), (xp, yp) = pair

        x_mean += x
        y_mean += y
        xp_mean += xp
        yp_mean += yp

    x_mean /= n
    y_mean /= n
    xp_mean /= n
    yp_mean /= n

    s_x_xp = 0
    s_y_yp = 0
    s_x_yp = 0
    s_y_xp = 0
    for pair in point_pairs:

        (x, y), (xp, yp) = pair

        s_x_xp += (x - x_mean)*(xp - xp_mean)
        s_y_yp += (y - y_mean)*(yp - yp_mean)
        s_x_yp += (x - x_mean)*(yp - yp_mean)
        s_y_xp += (y - y_mean)*(xp - xp_mean)

    rot_angle = math.atan2(s_x_yp - s_y_xp, s_x_xp + s_y_yp)
    translation_x = xp_mean - (x_mean*math.cos(rot_angle) - y_mean*math.sin(rot_angle))
    translation_y = yp_mean - (x_mean*math.sin(rot_angle) + y_mean*math.cos(rot_angle))

    return rot_angle, translation_x, translation_y

def icp(reference_points, points, max_iterations=100, distance_threshold=1000, convergence_translation_threshold=1e-3,
        convergence_rotation_threshold=1e-4, point_pairs_threshold=3, verbose=False):
    """
    An implementation of the Iterative Closest Point algorithm that matches a set of M 2D points to another set
    of N 2D (reference) points.

    :param reference_points: the reference point set as a numpy array (N x 2)
    :param points: the point that should be aligned to the reference_points set as a numpy array (M x 2)
    :param max_iterations: the maximum number of iteration to be executed
    :param distance_threshold: the distance threshold between two points in order to be considered as a pair
    :param convergence_translation_threshold: the threshold for the translation parameters (x and y) for the
                                              transformation to be considered converged
    :param convergence_rotation_threshold: the threshold for the rotation angle (in rad) for the transformation
                                               to be considered converged
    :param point_pairs_threshold: the minimum number of point pairs the should exist
    :param verbose: whether to print informative messages about the process (default: False)
    :return: the transformation history as a list of numpy arrays containing the rotation (R) and translation (T)
             transformation in each iteration in the format [R | T] and the aligned points as a numpy array M x 2
    """

    transformation_history = []

    nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(reference_points)

    for iter_num in range(max_iterations):
        if verbose:
            print('------ iteration', iter_num, '------')

        closest_point_pairs = []  # list of point correspondences for closest point rule

        distances, indices = nbrs.kneighbors(points)
        for nn_index in range(len(distances)):
            if distances[nn_index][0] < distance_threshold:
                closest_point_pairs.append((points[nn_index], reference_points[indices[nn_index][0]]))

        # if only few point pairs, stop process
        if verbose:
            print('number of pairs found:', len(closest_point_pairs))
        if len(closest_point_pairs) < point_pairs_threshold:
            if verbose:
                print('No better solution can be found (very few point pairs)!')
            break

        # compute translation and rotation using point correspondences
        closest_rot_angle, closest_translation_x, closest_translation_y = point_based_matching(closest_point_pairs)
        if closest_rot_angle is not None:
            if verbose:
                print('Rotation:', math.degrees(closest_rot_angle), 'degrees')
                print('Translation:', closest_translation_x, closest_translation_y)
        if closest_rot_angle is None or closest_translation_x is None or closest_translation_y is None:
            if verbose:
                print('No better solution can be found!')
            break

        # transform 'points' (using the calculated rotation and translation)
        c, s = math.cos(closest_rot_angle), math.sin(closest_rot_angle)
        rot = np.array([[c, -s],
                        [s, c]])
        aligned_points = np.dot(points, rot.T)
        aligned_points[:, 0] += closest_translation_x
        aligned_points[:, 1] += closest_translation_y

        # update 'points' for the next iteration
        points = aligned_points

        # update transformation history
        transformation_history.append(np.hstack((rot, np.array([[closest_translation_x], [closest_translation_y]]))))

        # check convergence
        if (abs(closest_rot_angle) < convergence_rotation_threshold) \
                and (abs(closest_translation_x) < convergence_translation_threshold) \
                and (abs(closest_translation_y) < convergence_translation_threshold):
            if verbose:
                print('Converged!')
            break

    return transformation_history, points

ext: py

max_stars_repo_name: vortexntnu/Vortex-CV

max_issues_repo_name: vortexntnu/Vortex-CV

reason: import numpy


text:
import os
import json
import pandas as pd
import argparse
import numpy as np
from pathlib import Path

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

import seaborn as sns

from .utils import *

plt.style.use('ggplot')

ILURL_HOME = os.environ['ILURL_HOME']

EMISSION_PATH = \
    f'{ILURL_HOME}/data/emissions'

FIGURE_X = 6.0
FIGURE_Y = 4.0

def get_arguments():
    parser = argparse.ArgumentParser(
        description="""
            This script evaluates a traffic light system.
        """
    )
    parser.add_argument('experiment_root_folder', type=str, nargs='?',
                        help='Experiment root folder.')

    return parser.parse_args()

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def print_arguments(args):

    print('Arguments (analysis/test_plots.py):')
    print('\tExperiment root folder: {0}\n'.format(args.experiment_root_folder))

if __name__ == "__main__":

    print('\nRUNNING analysis/metrics_per_route.py\n')

    args = get_arguments()
    print_arguments(args)
    experiment_root_folder = args.experiment_root_folder

    # Prepare output folder.
    output_folder_path = os.path.join(experiment_root_folder, 'plots/test')
    print('Output folder: {0}'.format(output_folder_path))
    os.makedirs(output_folder_path, exist_ok=True)

    # Get cycle length from tls_config.json file.
    config_files = list(Path(experiment_root_folder).rglob('tls_config.json'))
    with config_files[0].open('r') as f:
        json_file = json.load(f)
    cycle_time = json_file['rl']['cycle_time']

    # Get all *.csv files from experiment root folder.
    csv_files = [str(p) for p in list(Path(experiment_root_folder).rglob('*-emission.csv'))]

    print('Number of csv files found: {0}'.format(len(csv_files)))

    for csv_file in csv_files:

        print('Processing CSV file: {0}'.format(csv_file))

        # Load CSV data.
        df_csv = get_emissions(csv_file)

        df_per_vehicle = get_vehicles(df_csv)

        # Randomly pick some routes.
        all_routes = df_per_vehicle['route'].unique()
        routes = np.random.choice(all_routes, 6)

        """
            PER-ROUTE METRICS PLOTS: ALL DAY .
        """
        # Speed per route - all day.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
             sns.distplot(df_per_vehicle[df_per_vehicle['route'] == route]['speed'], hist=False, kde=True,
                 kde_kws = {'linewidth': 3}, label=route)

        # All routes.
        sns.distplot(df_per_vehicle['speed'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average speed (m/s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average speed for different routes\n(Default)')
        plt.savefig('{0}/test_speeds_hist_per_route_default.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        # Travel time per route.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[df_per_vehicle['route'] == route]['total'], hist=False, kde=True,
                 kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle['total'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average travel time (s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average travel time for different routes\n(Default)')
        plt.savefig('{0}/test_travel_time_hist_per_route_default.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        # Waiting time per route.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[df_per_vehicle['route'] == route]['waiting'], hist=False, kde=True,
                 kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle['waiting'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average waiting time (s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average waiting time for different routes\n(Default)')
        plt.savefig('{0}/test_waiting_time_hist_per_route_default.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        """
            PER-ROUTE METRICS PLOTS (FREE-FLOW ).
        """
        FREE_FLOW_TIME_CUT = 20000
        # Speed per route - free-flow.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[(df_per_vehicle['route'] == route) & (df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT)]['speed'], hist=False, kde=True,
                    kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle[(df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT)]['speed'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average speed (m/s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average speed for different routes\n(Free-flow)')
        plt.savefig('{0}/test_speeds_hist_per_route_free_flow.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        # Travel time per route - free-flow.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[(df_per_vehicle['route'] == route) & (df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT)]['total'], hist=False, kde=True,
                    kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle[(df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT)]['total'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average travel time (s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average travel time for different routes\n(Free-flow)')
        plt.savefig('{0}/test_travel_time_hist_per_route_free_flow.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        # Waiting time per route free-flow.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[(df_per_vehicle['route'] == route) & (df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT)]['waiting'], hist=False, kde=True,
                    kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle[(df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT)]['waiting'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average waiting time (s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average waiting time for different routes\n(Free-flow)')
        plt.savefig('{0}/test_waiting_time_hist_per_route_free_flow.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        """
            PER-ROUTE METRICS PLOTS (CONGESTED).
        """
        CONGESTED_TIME_LOW_CUT = 25000
        CONGESTED_TIME_HIGH_CUT = 40000
        # Speed per route congested.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[(df_per_vehicle['route'] == route) & (df_per_vehicle['finish'] < CONGESTED_TIME_HIGH_CUT) & (df_per_vehicle['finish'] > CONGESTED_TIME_LOW_CUT)]['speed'], hist=False, kde=True,
                    kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle[(df_per_vehicle['finish'] < CONGESTED_TIME_HIGH_CUT) & (df_per_vehicle['finish'] > CONGESTED_TIME_LOW_CUT)]['speed'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average Speed (m/s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average speed for different routes\n(Congested)')
        plt.savefig('{0}/test_speeds_hist_per_route_congested.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        # Travel time per route congested.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:

            sns.distplot(df_per_vehicle[(df_per_vehicle['route'] == route) & (df_per_vehicle['finish'] < CONGESTED_TIME_HIGH_CUT) & (df_per_vehicle['finish'] > CONGESTED_TIME_LOW_CUT)]['total'], hist=False, kde=True,
                    kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle[(df_per_vehicle['finish'] < CONGESTED_TIME_HIGH_CUT) & (df_per_vehicle['finish'] > CONGESTED_TIME_LOW_CUT)]['total'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average travel time (s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average travel time for different routes\n(Congested)')
        plt.savefig('{0}/test_travel_time_hist_per_route_congested.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

        # Waiting time per route congested.
        fig = plt.figure()
        fig.set_size_inches(FIGURE_X, FIGURE_Y)

        for route in routes:
            sns.distplot(df_per_vehicle[(df_per_vehicle['route'] == route) & (df_per_vehicle['finish'] < CONGESTED_TIME_HIGH_CUT) & (df_per_vehicle['finish'] > CONGESTED_TIME_LOW_CUT)]['waiting'], hist=False, kde=True,
                    kde_kws = {'linewidth': 3}, label=route)

        sns.distplot(df_per_vehicle[(df_per_vehicle['finish'] < FREE_FLOW_TIME_CUT) & (df_per_vehicle['finish'] > CONGESTED_TIME_LOW_CUT)]['waiting'], hist=False, kde=True, color='black',
                 kde_kws = {'linewidth': 3, 'linestyle': '--'}, label='All routes')

        plt.legend()
        plt.xlabel('Average waiting time (s)')
        plt.ylabel('Density')
        plt.title('Vehicles\' average waiting time for different routes\n(Congested)')
        plt.savefig('{0}/test_waiting_time_hist_per_route_congested.png'.format(output_folder_path), bbox_inches='tight', pad_inches=0)
        plt.close()

max_stars_repo_name: CubeSkyy/ILU-RL

max_issues_repo_name: CubeSkyy/ILU-RL

max_forks_repo_name: CubeSkyy/ILU-RL


text:
from __future__ import division
from matplotlib.animation import FuncAnimation
import gsd.fl
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
from hoomd.data import boxdim
from scipy.spatial import Delaunay

def __init__():
    """
    This package is supposed to operate on data calculated in a rectangular box.
    The package assumes this and doesn't check.
    """
    pass

def correct_jumps(pos, pos_m1_cntn, pos_m1_inbox, lx, ly):
    """
    Correct jumps over the periodic boundary. Jumps by lx or ly make diffusion calculation
    nonsense.
    pos - positions in current frame
    pos_m1_cntc - "continuous" position in the previous frame: corrected for jumps from previous step
    pos_m1_inbox - "raw" position in the previous frame: inside simulation box
    lx, ly - unit cell length and width

    return: corrected positions in current frame
    """

    dr = pos - pos_m1_inbox
    # x-direction jumps:
    ind = np.where(dr[:,0] > lx/2)
    dr[ind, 0] = dr[ind, 0] - lx
    ind = np.where(dr[:,0] < - lx/2)
    dr[ind, 0] = dr[ind, 0] + lx

    #y-jumps:
    ind = np.where(dr[:,1] > ly/2)
    dr[ind, 1] = dr[ind, 1] - ly
    ind = np.where(dr[:,1] < - ly/2)
    dr[ind, 1] = dr[ind, 1] + ly

    return pos_m1_cntn + dr

def get_file_list(folder_path):
    """
        Read filenames and corresponding Gamma and T_values in a folder
    """
    f_list, gamma_list, T_list, dt_list = [], [], [], []
    if folder_path[-1] != '/':
        folder_path = folder_path + '/'
    try:
        fl = open(folder_path + 'list.txt', 'r')
    except:
        print('Could not find list.txt in {}'.format(folder_path))
        raise FileNotFoundError
    for line in fl.readlines():
        if line[0] != '#':
            words = line.split('\t')
            f_list.append(words[0])
            gamma_list.append(float(words[1]))
            T_list.append(float(words[2]))
            dt_list.append(float(words[3]))
    fl.close()
    return f_list, gamma_list, T_list, dt_list

def diffusion_from_gsd(folder_path, center_fixed = True, useframes = -1):
    """
    Calculate diffusion coefficients vs Gamma from gsd files located in folder_path.
    The folder must have list.txt that has the following structure:

    # file    Gamma    T    dt
    00000.gsd	145.00000000	1.09746597	0.00100000
    00001.gsd	144.28571429	1.10289897	0.00100000

    Lines starting with # are ignored. The columns are tab-separated
    Each gsd file should contain a long enough trajectory in thermalized state.

    Diffusion constant D is calculated from 4Dt = <(r(t) - r(0))^2>.
    The average is calculated over all particles and over different time origins.
    Time origins go from 0 to n_frames/2, and t goes from 0 to n_frames/2. This way,
    the data are always within the trajectory.

    center_fixed = True: eliminate oveall motion of center of mass

    return D_x, D_y, gamma_list, T_list
    D_x, D_y diffusion for x- and y-coordinates;
    all arrays are ordered for right correspondence, e.g. D_x[i] <-> gamma_list[i]
    """
    f_list, gamma_list, T_list, dt_list = get_file_list(folder_path)

    D_x_list = np.zeros(len(f_list))
    D_y_list = np.zeros(len(f_list))

    if folder_path[-1] != '/':
        folder_path = folder_path + '/'

    for i, f_name in enumerate(f_list):
        print(f_name)
        with gsd.fl.GSDFile(folder_path + f_name, 'rb') as f:
            n_frames = f.nframes
            box = f.read_chunk(frame=0, name='configuration/box')
            half_frames = int(n_frames/2) - 1 #sligtly less than half to avoid out of bound i
            if useframes < 1 or useframes > half_frames:
                useframes = half_frames
            t_step0 = f.read_chunk(frame=0, name='configuration/step')
            t_step1 = f.read_chunk(frame=1, name='configuration/step')
            snap_period = t_step1[0] - t_step0[0]
            n_p = f.read_chunk(frame=0, name='particles/N')
            if i == 0: #create square-average displacement once
                x_sq_av = np.zeros((useframes, len(f_list)))
                y_sq_av = np.zeros((useframes, len(f_list)))
            for t_origin in range(n_frames - useframes - 1):
                pos_0 = f.read_chunk(frame=t_origin, name='particles/position')
                mean_pos_0 = np.mean(pos_0, axis = 0)
                pos = pos_0
                pos_raw = pos_0
                for j_frame in range(useframes):
                    pos_m1 = pos
                    pos_m1_raw = pos_raw
                    pos_raw = f.read_chunk(frame=j_frame + t_origin, name='particles/position') - pos_0
                    pos = correct_jumps(pos_raw, pos_m1, pos_m1_raw, box[0], box[1])
                    if center_fixed:
                        pos -= np.mean(pos, axis = 0) - mean_pos_0 #correct for center of mass movement

                    x_sq_av[j_frame, i] += np.mean(pos[:,0]**2)
                    y_sq_av[j_frame, i] += np.mean(pos[:,1]**2)
            x_sq_av[:, i] /= (n_frames - useframes - 1)
            y_sq_av[:, i] /= (n_frames - useframes - 1)
            # OLS estimate for beta_x[0] + beta_x[1]*t = <|x_i(t) - x_i(0)|^2>
            a = np.ones((useframes, 2)) # matrix a = ones(half_frames) | (0; dt; 2dt; 3dt; ...)
            a[:,1] = snap_period*dt_list[i]*np.cumsum(np.ones(useframes), axis = 0) - dt_list[i]
            b_cutoff = int(useframes/10) #cutoff to get only linear part of x_sq_av, makes results a bit more clean
            beta_x = np.linalg.lstsq(a[b_cutoff:, :], x_sq_av[b_cutoff:,i], rcond=-1)
            beta_y = np.linalg.lstsq(a[b_cutoff:, :], y_sq_av[b_cutoff:,i], rcond=-1)
            D_x_list[i] = beta_x[0][1]/4
            D_y_list[i] = beta_y[0][1]/4
    return D_x_list, D_y_list, gamma_list, T_list

def gsd_trajectory(fpath, axis, periodic = False, center_fixed = True):
    """
    Return a trajectory data for all particles in gsd file located at fpath;
    axis = 0 for x coordinates, 1 for y coordinates;
    periodic = True to contain particles within box, False - allow particles to diffuse out of box;

    """
    if axis != 0 and axis != 1:
        raise ValueError('axis must be 0 or 1')

    with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
        n_frames = f_gsd.nframes
        n_p = f_gsd.read_chunk(frame=0, name='particles/N')
        box = f_gsd.read_chunk(frame=0, name='configuration/box')
        pos = np.zeros((n_frames, n_p[0]))
        pos_frame = f_gsd.read_chunk(frame=0, name='particles/position')
        mean_pos_0 = np.mean(pos_frame, axis = 0)
        for j_frame in range(n_frames):
            pos_m1 = pos_frame
            pos_frame = f_gsd.read_chunk(frame=j_frame, name='particles/position')
            if not periodic:
                pos_frame = correct_jumps(pos_frame, pos_m1, box[0], box[1])
            if center_fixed:
                pos_frame -= np.mean(pos_frame, axis = 0) - mean_pos_0
            pos[j_frame, :] = pos_frame[:, axis]
    return pos

def find_neighbors(pos, box, rcut = 1.4):
    """ Find neighbors for each particle. The neighbors are determined to be within rcut distance from the
    particle.
    \param pos - N x 3 array of positions
    \param box - hoomd box object, simulation box
    \param rcut - cutoff radius for neighbors
    Return neighbor_list and neighbor_num.
    neighbor_list - N x 30 array of int, each row i containing indices of neighbors of i-th particle,
                    including the particle itself.
                    The remaining part of the row is filled with -1 (e.g. for 6 neighbors, remaining 23 sites are -1).
    neighbor_num - int array of size N containing numbers of neighbors for each particle.
    """
    neighbor_list = np.zeros((pos.shape[0], 30), dtype=int) - 1
    neighbor_num = np.zeros(pos.shape[0], dtype = int)
    dist_list = np.zeros((pos.shape[0], 30), dtype=float) - 1

    for i, r in enumerate(pos):
        pos_ref = reshape_to_box(pos, r, box)
        box_ind = np.where((pos_ref[:,0] > -rcut)*(pos_ref[:,0] < rcut)*(pos_ref[:,1] > -rcut)*(pos_ref[:,1] < rcut))[0]
        box_pos = pos_ref[box_ind, :]
        dist = np.sqrt(np.sum(box_pos**2, axis = 1))
        box_neighbors = np.where(dist < rcut)[0]
        box_dist = dist[box_neighbors]
        neighbor_ind = box_ind[box_neighbors]
        neighbor_num[i] = len(neighbor_ind) - 1
        for j, ind in enumerate(neighbor_ind):
            neighbor_list[i, j] = ind
            dist_list[i, j] = box_dist[j]
    return neighbor_list, neighbor_num

def create_virtual_layer(pos, box):
    """
    Create a virtual layer around the simulation box of ~2 interelectron distances of particles that are periodic replicas from the
    simulation box.
    """
    N = pos.shape[0]
    a = 2*np.sqrt(box.Lx*box.Ly/N) # approximately 1-3 interelectron distances
    # Add full periodic replicas around the box
    shift_arr = [-1, 0, 1]
    full_virtual_pos = np.zeros((N*9, 3))
    counter = 0
    for i in shift_arr:
        for j in shift_arr:
            #i,j = 0,0 when counter=4
            pos_shifted = np.zeros(pos.shape)
            pos_shifted[:,0] = pos[:,0] + i*box.Lx
            pos_shifted[:,1] = pos[:,1] + j*box.Ly
            full_virtual_pos[counter*N:(counter + 1)*N, :] = pos_shifted
            counter += 1
    # Crop the extended virtual box to the original box plus a layer of width `a`
    box_ind = np.where((np.abs(full_virtual_pos[:,0]) < 0.5*box.Lx + a)*(np.abs(full_virtual_pos[:,1]) < 0.5*box.Ly + a))[0]
    crop_virtual_pos = full_virtual_pos[box_ind, :]

    return crop_virtual_pos, box_ind

def find_neighbors_delone(pos, box):
    """ Find neighbors for each particle with the Delone triangulation methos.
    \param pos - N x 3 array of positions
    \param box - hoomd box object, simulation box
    \param rcut - cutoff radius for neighbors
    Return neighbor_list and neighbor_num.
    neighbor_list - N x 30 array of int, each row i containing indices of neighbors of i-th particle,
                    including the particle itself.
                    The remaining part of the row is filled with -1 (e.g. for 6 neighbors, remaining 23 sites are -1).
    neighbor_num - int array of size N containing numbers of neighbors for each particle.
    """
    N = pos.shape[0]
    neighbor_list = np.zeros((pos.shape[0], 30), dtype=int) - 1
    neighbor_num = np.zeros(pos.shape[0], dtype = int)
    dist_list = np.zeros((pos.shape[0], 30), dtype=float) - 1
    virtual_pos, virtual_ind = create_virtual_layer(pos, box)
    tri = Delaunay(virtual_pos[:, 0:2])
    (indices, indptr) = tri.vertex_neighbor_vertices
    for vi in range(virtual_pos.shape[0]):
        i = virtual_ind[vi] #index of a particle in the big virtual box with full repeated cells
        i_real = i%N
        # for real particles in the real box
        if i >= N*4 and i < N*5:
            # indices of neighbors in the vitrual box
            vi_neighb = indptr[indices[vi]:indices[vi+1]]
            # recover real neighbor indices in the original box
            i_neighb = virtual_ind[vi_neighb]%N
            n_neighb = len(i_neighb)
            neighbor_list[i_real, :n_neighb] = i_neighb
            neighbor_num[i_real] = n_neighb
    return neighbor_list, neighbor_num

def animate_gsd(fpath, savefile = None, periodic = False, center_fixed = True, interval = 100, figsize = (12, 12), rcut = 1.4,\
               neighb = False):
    """
    Create animation from a gsd file, where fpath is the path+filename to the gsd file.
    lx, ly - dimensions of the rectangular simulation box;
    savefile - filename for saving the animation. Not trying to save if the value is None;
    periodic - show with periodic boundary conditions if True, let particles migrate out of box if False;
    interval - time interval between frames in microseconds;

    return fig, ax
    """
    def init():
        scat.set_offsets(pos[0, :, :])

        time_text.set_text('')
        return scat, time_text

    def update(frame_number):
        scat.set_offsets(pos[frame_number, :, :])
        time_text.set_text('frame = {} of {}'.format(frame_number, n_frames))
        if neighb:
            five_ind = np.where(neighbor_num[frame_number, :] == 5)[0]
            seven_ind = np.where(neighbor_num[frame_number, :] == 7)[0]
            five_scat.set_offsets(pos[frame_number, five_ind, :])
            seven_scat.set_offsets(pos[frame_number, seven_ind, :])
        else:
            five_scat.set_offsets(empty_pos)
            seven_scat.set_offsets(empty_pos)
        return scat, time_text, five_scat, seven_scat

    with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
        n_frames = f_gsd.nframes
        n_p = f_gsd.read_chunk(frame=0, name='particles/N')
        box = f_gsd.read_chunk(frame=0, name='configuration/box')
        pos = np.zeros((n_frames, n_p[0], 2))
        neighbor_num = np.zeros((n_frames, int(n_p)), dtype = int)
        pos_frame = f_gsd.read_chunk(frame=0, name='particles/position')
        pos_frame_raw = pos_frame
        mean_pos_0 = np.mean(pos_frame, axis = 0)
        for j_frame in range(n_frames):
            pos_m1 = pos_frame
            pos_m1_raw = pos_frame_raw
            pos_frame_raw = f_gsd.read_chunk(frame=j_frame, name='particles/position')
            if not periodic:
                pos_frame = correct_jumps(pos_frame_raw, pos_m1, pos_m1_raw, box[0], box[1])
            else:
                pos_frame = pos_frame_raw
            if center_fixed:
                pos_frame -= np.mean(pos_frame, axis = 0) - mean_pos_0
            pos[j_frame, :, :] = pos_frame[:, 0:2]
            if neighb:
                boxdim_box = boxdim(box[0], box[1], box[2])
                neighbor_list, neighbor_num[j_frame, :] = find_neighbors_delone(pos[j_frame, :,:], boxdim_box, rcut = rcut)

    fig = plt.figure(figsize = figsize)
    ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                         xlim=(-box[0], box[0]), ylim=(-box[1], box[1]))

    scat = ax.scatter(pos[0, :, 0], pos[0, :, 1],
                      s = 3,
                      facecolors='blue')
    empty_pos = np.zeros(0)
    if neighb:
        five_ind_0 = np.where(neighbor_num[0, :] == 5)[0]
        seven_ind_0 = np.where(neighbor_num[0, :] == 7)[0]
        seven_scat = ax.scatter(pos[0, seven_ind_0, 0], pos[0, seven_ind_0, 1],
                          s = 5,
                          facecolors='green')
        five_scat = ax.scatter(pos[0, five_ind_0, 0], pos[0, five_ind_0, 1],
                          s = 5,
                          facecolors='red')
    else:
        seven_scat = ax.scatter(empty_pos, empty_pos)
        five_scat = ax.scatter(empty_pos, empty_pos)
    time_text = ax.text(0.02, 1.05, '', transform=ax.transAxes)

    animation = FuncAnimation(fig, update, interval=interval, frames=n_frames, blit=True)

    if not savefile == None:
        try:
            animation.save(savefile, fps=30)
        except Exception as ee: print(ee)
    return fig, ax, animation

def plot_DxDy(Dx, Dy, gamma_list, timestamp, text_list = [], text_pos = 'c', folder = ''):
    """Create two-panel plot for Dx, Dy vs Gamma
    text_list = list of strings that will be displayed on the plot
    text_pos = 'l', 'c', or 'r' for left, center, and right alignment of the text
    timestamp = string that will serve as a base for the filename
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize = (9, 6))
    ax1.scatter(gamma_list, Dx)
    ax1.set_xlabel('$\Gamma$')
    ax1.set_ylabel('$D_x$')

    ax2.scatter(gamma_list, Dy)
    ax2.set_xlabel('$\Gamma$')
    ax2.set_ylabel('$D_y$')
    fig.tight_layout()
    fig.patch.set_alpha(1)

    #determine text position
    y_lim = ax1.get_ylim()
    x_lim = ax1.get_xlim()
    h = y_lim[1] - y_lim[0]
    w = x_lim[1] - x_lim[0]
    text_y = y_lim[1] - 0.1*h

    text_x = {
        'l': x_lim[0] + 0.1*w,
        'c': x_lim[0] + 0.4*w,
        'r': x_lim[0] + 0.7*w
    }.get(text_pos, x_lim[0] + 0.4*w)

    #print text
    if type(text_list) == list:
        n_str = len(text_list)
        for i in range(n_str):
            ax1.text(text_x, text_y - 0.05*h*i, text_list[i])
    elif type(text_list) == str:
        ax1.text(text_x, text_y, text_list)
    else:
        raise TypeError('text_list must be a list of strings or a string')
    if folder != '':
        if folder[-1] != '/':
            folder = folder + '/'
    fig.savefig(folder + timestamp + '_diff.png')
    return fig, ax1, ax2

def plot_positions(system=None, pos=None, box=None, figsize = (7, 7), gridon = True, ax=None, fig=None, s=5):
    """ Show positions of all particles in the system,
    where system is the hoomd system, produced by hoomd.init.
    Show grid lines if gridon == True
    """
    if system != None:
        snapshot = system.take_snapshot(all=True)
        box = snapshot.box
        pos = snapshot.particles.position
    if ax==None or fig==None:
        fig = plt.figure(figsize = figsize)
        ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                             xlim=(-0.6*box.Lx, 0.6*box.Lx), ylim=(-0.6*box.Ly, 0.6*box.Ly))

    scat = ax.scatter(pos[:, 0], pos[:, 1],
                      s = s,
                      facecolors='blue')
    if gridon:
        ax.grid()
    return fig, ax

def reshape_to_box(pos, ref, box):
    """
    Center the simulation box arund ref, and return particles to within the simulation box.
    Return new positions.
    """
    pos_ref = pos - ref
    mask = np.fix(pos_ref[:, 0]/(0.5*box.Lx)) #-1 or 1 if to the left or right from the centered box
    pos_ref[:, 0] -= mask*box.Lx

    mask = np.fix(pos_ref[:, 1]/(0.5*box.Ly))
    pos_ref[:, 1] -= mask*box.Ly
    return pos_ref

def is_odd(num):
    return num & 0x1

def pair_correlation(pos, box, n_bins = (100, 100)):
    """ Calculate pair correlation function for a given snapshot.
        \param positions: N x 3 array, particle positions
        \param box: simulation box object with fields Lx and Ly
        \param n_bins: tuple of size 2, numbers of bins in x and y direction for correlation function, both even
        Return: array of shape n_bins, pair correlation function normalized to (N - 1). Zero of coordinate is in the middle
        of the pixel grid (between pixels indexed (n_bins[i]/2 - 1) and (n_bins[i]/2)).
    """
    try:
        if is_odd(n_bins[0]) or is_odd(n_bins[1]):
            raise ValueError("n_bins must be 2 x 1 tuple of even numbers")
    except:
        raise ValueError("n_bins must be 2 x 1 tuple of even numbers")

    bins_x = np.linspace(-box.Lx/2, box.Lx/2, n_bins[0] + 1)
    bins_y = np.linspace(-box.Ly/2, box.Ly/2, n_bins[1] + 1)
    g = np.zeros(n_bins)
    for r in pos:
        pos_ref = reshape_to_box(pos, r, box)
        ind_x = np.digitize(pos_ref[:,0], bins_x) - 1
        ind_y = np.digitize(pos_ref[:,1], bins_y) - 1
        if np.max(ind_x) >= n_bins[0]:
            ind_x[np.where(ind_x >= n_bins[0])] = n_bins[0] - 1
        if np.max(ind_y) >= n_bins[1]:
            ind_y[np.where(ind_y >= n_bins[1])] = n_bins[1] - 1

        g[ind_x[:], ind_y[:]] += 1
    g[int(n_bins[0]/2), int(n_bins[1]/2)] = 0
    #Normalize so that the sum over all pixels is N - 1
    g /= pos.shape[0]
    #Normalize to the particle density so that integral of g is (N - 1)/n_s
    n_s = pos.shape[0]/box.Lx/box.Ly
    g *= n_bins[0]*n_bins[1]/box.Lx/box.Ly
    g /= n_s
    return g

def pair_correlation_from_gsd(filename, n_bins = (100, 100), frames =(0, -1)):
    """ Calculate pair correlation function.
        Averaged over frames in the range from frames[0] to frames[1]. Negative frames are counted from the end of time.
        \param filename: name of the gsd file
        \param n_bins: tuple of size 2, numbers of bins in x and y direction for correlation function, both even
        Return: array of shape n_bins, pair correlation function normalized to (N - 1). Zero of coordinate is in the middle
        of the pixel grid (between pixels indexed (n_bins[i]/2 - 1) and (n_bins[i]/2)).
    """
    try:
        if is_odd(n_bins[0]) or is_odd(n_bins[1]):
            raise ValueError("n_bins must be 2 x 1 tuple of even numbers")
    except:
        raise ValueError("n_bins must be 2 x 1 tuple of even numbers")

    g = np.zeros(n_bins)
    with gsd.fl.GSDFile(filename, 'rb') as f_gsd:
        n_frames_total = f_gsd.nframes
        if frames[0] > n_frames_total or frames[1] > n_frames_total:
            raise ValueError('frames beyond n_frames_total')
        #translate negative indices into positive domain:
        abs_frames = (frames[0] -(frames[0]//n_frames_total)*n_frames_total, \
                     frames[1] -(frames[1]//n_frames_total)*n_frames_total)
        if abs_frames[0] > abs_frames[1]:
            raise ValueError('frames[0] must be to the left from frames[1]')
        all_frames = np.arange(0, n_frames_total, 1)
        selected_frames = all_frames[abs_frames[0]:abs_frames[1] + 1]
        n_frames = abs_frames[1] - abs_frames[0] + 1

        n_p = f_gsd.read_chunk(frame=0, name='particles/N')
        box_array = f_gsd.read_chunk(frame=0, name='configuration/box')
        box = boxdim(*box_array[0:3])
        pos = np.zeros((n_frames, n_p[0], 2))
        for j_frame in range(n_frames):
            pos_frame = f_gsd.read_chunk(frame=selected_frames[j_frame], name='particles/position')
            g += pair_correlation(pos_frame, box, n_bins = n_bins)
    g /= n_frames
    return g

def plot_pair_correlation(g, box, figsize = (8,8), cmap = "plasma", interp = 'none', tickfont=18,\
                          nbins=50, alpha=1, fig=None, ax=None, origin_marker=True, minmax=None, imshow=False):
    if ax==None or fig==None:
        fig = plt.figure(figsize = figsize)
        ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                             xlim=(-0.5*box.Lx, 0.5*box.Lx), ylim=(-0.5*box.Ly, 0.5*box.Ly))
    x = np.linspace(-box.Lx/2, box.Lx/2, g.shape[0])
    y = np.linspace(-box.Ly/2, box.Ly/2, g.shape[1])
    X, Y = np.meshgrid(x, y)
    if minmax == None:
        levels = MaxNLocator(nbins=nbins).tick_values(g.min(), g.max())
        ticks = [g.min(), g.max()]
    else:
        levels = MaxNLocator(nbins=nbins).tick_values(minmax[0], minmax[1])
        ticks = minmax
    cmapo = plt.get_cmap(cmap)
    norm = BoundaryNorm(levels, ncolors=cmapo.N, clip=True)#

    #pc = ax.pcolor(X, Y, np.transpose(g), cmap = cmap,\
    #alpha=alpha, norm=norm, edgecolor='none', antialiaseds=False)
    if imshow:
        pc = ax.imshow(np.flipud(np.transpose(g)), extent=[X.min(), X.max(), Y.min(), Y.max()], cmap=cmap,\
                       alpha=alpha, interpolation=interp, norm=norm)
    else:
        pc = ax.pcolor(X, Y, np.transpose(g), cmap = cmapo,\
        alpha=alpha, norm=norm, edgecolor=None, antialiaseds=False, linewidth=0)
        pc.set_rasterized(True)
    if origin_marker:
        ax.scatter(0,0, c='r', marker = '+')
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    cbar = fig.colorbar(pc, cax=cax, ax=ax, ticks=ticks)#, fraction=0.056, pad=0.04)
    cbar.ax.tick_params(labelsize=tickfont)
    cbar.ax.set_yticklabels(['{:.2f}'.format(ticks[0]), '{:.2f}'.format(ticks[1])])

    return fig, ax, cax

def psi_order(pos, box, nx=100, ny=100, rcut=1.4):
    """
    \param pos - N x 3 array of positions
    \param box - hoomd box object, simulation box
    \param rcut - cutoff radius for neighbors
    Return neighbor_list and neighbor_num.
    neighbor_list - N x 30 array of int, each row i containing indices of neighbors of i-th particle,
                    including the particle itself.
                    The remaining part of the row is filled with -1 (e.g. for 6 neighbors, remaining 23 sites are -1).
    neighbor_num - int array of size N containing numbers of neighbors for each particle.
    """
    neighbor_list = np.zeros((pos.shape[0], 30), dtype=int) - 1
    neighbor_num = np.zeros(pos.shape[0], dtype = int)
    dist_list = np.zeros((pos.shape[0], 30), dtype=float) - 1
    psi = np.zeros((ny, nx), dtype=complex)
    hx = box.Lx/nx
    hy = box.Ly/ny
    X = np.linspace(-0.5*box.Lx, 0.5*box.Lx - hx, nx)
    Y = np.linspace(-0.5*box.Ly, 0.5*box.Ly - hy, ny)
    r = np.zeros(3)

    for i in range(nx):
        for j in range(ny):
            r[0] = X[i]
            r[1] = Y[j]
            pos_ref = reshape_to_box(pos, r, box)
            r_ind = np.where((pos_ref[:,0] > -rcut)*(pos_ref[:,0] < rcut)*(pos_ref[:,1] > -rcut)*(pos_ref[:,1] < rcut))[0]
            r_pos = pos_ref[r_ind, :]
            dist = np.sqrt(np.sum(r_pos**2, axis = 1))
            r_nearest_ind = np.argmin(dist)
            nearest_ind = r_ind[r_nearest_ind]
            nr = pos_ref[nearest_ind, :] # coordinate of nearest particle in r-centered box
            box_ind = np.where((pos_ref[:,0] > -rcut + nr[0])*(pos_ref[:,0] < rcut + nr[0])*\
                               (pos_ref[:,1] > -rcut + nr[1])*(pos_ref[:,1] < rcut + nr[1]))[0]
            box_pos = pos_ref[box_ind, :]
            dist = np.sqrt(np.sum((box_pos - nr)**2, axis = 1))
            box_neighbors = np.where(dist < rcut)[0]
            nb_rv = box_pos[box_neighbors, :] - nr # radius-vectors to nearest neighbors
            nb_dist = np.sqrt(np.sum(nb_rv**2, axis = 1))
            nb_rv = nb_rv[np.where(nb_dist != 0)]
            nb_dist = nb_dist[np.where(nb_dist != 0)]
            cos_theta = nb_rv[:,0]/nb_dist
            theta = np.arccos(cos_theta) # sign of the angle still uncertain
            neg_y_ind = np.where(nb_rv[:,1] < 0)[0]
            theta[neg_y_ind] = 2*np.pi - theta[neg_y_ind] #resolve uncertainty from arccos
            psi[j,i] = np.mean(np.exp(1j*6*theta))
    return psi

def psi_order_delone(pos, box, nx=100, ny=100, rcut=1.4):
    """
    \param pos - N x 3 array of positions
    \param box - hoomd box object, simulation box
    \param rcut - cutoff radius for neighbors
    Return neighbor_list and neighbor_num.
    neighbor_list - N x 30 array of int, each row i containing indices of neighbors of i-th particle,
                    including the particle itself.
                    The remaining part of the row is filled with -1 (e.g. for 6 neighbors, remaining 23 sites are -1).
    neighbor_num - int array of size N containing numbers of neighbors for each particle.
    """
    #neighbor_list = np.zeros((pos.shape[0], 30), dtype=int) - 1
    #neighbor_num = np.zeros(pos.shape[0], dtype = int)
    dist_list = np.zeros((pos.shape[0], 30), dtype=float) - 1
    psi = np.zeros((ny, nx), dtype=complex)
    hx = box.Lx/nx
    hy = box.Ly/ny
    X = np.linspace(-0.5*box.Lx, 0.5*box.Lx - hx, nx)
    Y = np.linspace(-0.5*box.Ly, 0.5*box.Ly - hy, ny)
    r = np.zeros(3)
    n_list, n_num = find_neighbors_delone(pos, box)

    for i in range(nx):
        for j in range(ny):
            r[0] = X[i]
            r[1] = Y[j]
            pos_ref = reshape_to_box(pos, r, box)
            r_ind = np.where((pos_ref[:,0] > -rcut)*(pos_ref[:,0] < rcut)*(pos_ref[:,1] > -rcut)*(pos_ref[:,1] < rcut))[0]
            r_pos = pos_ref[r_ind, :]
            dist = np.sqrt(np.sum(r_pos**2, axis = 1))
            r_nearest_ind = np.argmin(dist)
            nearest_ind = r_ind[r_nearest_ind]
            nr = pos_ref[nearest_ind, :] # coordinate of nearest particle in r-centered box
            #box_ind = np.where((pos_ref[:,0] > -rcut + nr[0])*(pos_ref[:,0] < rcut + nr[0])*\
            #                   (pos_ref[:,1] > -rcut + nr[1])*(pos_ref[:,1] < rcut + nr[1]))[0]
            #box_pos = pos_ref[box_ind, :]
            #dist = np.sqrt(np.sum((box_pos - nr)**2, axis = 1))
            #box_neighbors = np.where(dist < rcut)[0]
            neighbors_ind = n_list[nearest_ind, :n_num[nearest_ind]]
            nb_rv = pos_ref[neighbors_ind, :] - nr # radius-vectors to nearest neighbors
            nb_dist = np.sqrt(np.sum(nb_rv**2, axis = 1))
            nb_rv = nb_rv[np.where(nb_dist != 0)]
            nb_dist = nb_dist[np.where(nb_dist != 0)]
            cos_theta = nb_rv[:,0]/nb_dist
            theta = np.arccos(cos_theta) # sign of the angle still uncertain
            neg_y_ind = np.where(nb_rv[:,1] < 0)[0]
            theta[neg_y_ind] = 2*np.pi - theta[neg_y_ind] #resolve uncertainty from arccos
            psi[j,i] = np.mean(np.exp(1j*6*theta))
    return psi

def psi6_order_from_gsd(fpath, frame=0):
    with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
        n_frames_total = f_gsd.nframes
        if frame > n_frames_total:
            raise ValueError('frames beyond n_frames_total')
        #translate negative indices into positive domain:
        abs_frame = frame -(frame//n_frames_total)*n_frames_total
        pos = f_gsd.read_chunk(frame=abs_frame, name='particles/position')
        box_array = f_gsd.read_chunk(frame=0, name='configuration/box')
        box = boxdim(*box_array[0:3])
    psi = psi_order_delone(pos, box, nx=100, ny=100, rcut=1.3)
    return psi

def compute_psi6_correlation_from_gsd(fpath, Nframes=1, frame_step=1, nxny=(100,100)):
    nx, ny = nxny
    cf_psi = np.zeros((ny, nx), dtype=complex)
    for frame in range(0, Nframes, frame_step):
        with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
            n_frames_total = f_gsd.nframes
            if frame > n_frames_total:
                print('frame {} beyond n_frames_total = {}'.format(frame, n_frames_total))
            else:
                #translate negative indices into positive domain:
                abs_frame = frame -(frame//n_frames_total)*n_frames_total
                pos = f_gsd.read_chunk(frame=abs_frame, name='particles/position')
                box_array = f_gsd.read_chunk(frame=0, name='configuration/box')
                box = boxdim(*box_array[0:3])
                psi = psi_order(pos, box, nx=nx, ny=ny, rcut=1.3)
                cf_psi += correlation_function(psi)
    cf_psi /= min(Nframes, n_frames_total)
    return cf_psi

def plot_delone_triangulation(fpath, frame=0, fig=None, ax=None):
    with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
        n_frames_total = f_gsd.nframes
        if frame > n_frames_total:
            raise ValueError('frames beyond n_frames_total')
        #translate negative indices into positive domain:
        abs_frame = frame -(frame//n_frames_total)*n_frames_total
        pos = f_gsd.read_chunk(frame=abs_frame, name='particles/position')
        box_array = f_gsd.read_chunk(frame=0, name='configuration/box')
        box = boxdim(*box_array[0:3])

    if ax==None or fig==None:
        fig = plt.figure(figsize = figsize)
        ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,
                             xlim=(-0.6*box.Lx, 0.6*box.Lx), ylim=(-0.6*box.Ly, 0.6*box.Ly))
    xlim = (-0.5*box.Lx, 0.5*box.Lx)
    ylim = (-0.5*box.Ly, 0.5*box.Ly)

    virtual_pos, virtual_ind = create_virtual_layer(pos, box)
    tri = Delaunay(virtual_pos[:, 0:2])
    ax.triplot(virtual_pos[:,0], virtual_pos[:,1], tri.simplices.copy(), color='black')
    ax.set_xlim(xlim)
    ax.set_ylim(ylim)
    return fig, ax

def correlation_function(f):
    cf = 0*f
    nynx = cf.shape
    #print('================================================================')
    for j in range(-nynx[0]//2, nynx[0]//2):
        for i in range(-nynx[1]//2, nynx[1]//2):
            #print('i={}, j={}, coor_i={}, coor_j={}'.format(i,j, nynx[1]//2 + i, nynx[0]//2 + j))
            f1 = np.roll(f, -j, axis=0)
            f2 = np.roll(f1, -i, axis=1)
            cf[nynx[0]//2 + j,nynx[1]//2 + i] = np.mean(f*np.conj(f2))
            #print('cf={}'.format(cf[nynx[0]//2 + j,nynx[1]//2 + i]))
    return cf

def Eee_from_gsd(fpath, table_path, width, height, step = 1):
    """
    Return interaction energy along a trajectory in gsd file located at fpath;
    table - interaction pair force object
    """
    import hoomd
    import hoomd.md
    hoomd.context.initialize('--mode=cpu')
    #hoomd = imp.reload(hoomd)
    system = hoomd.init.read_gsd(fpath, frame = 0)
    dt = 0.001
    all = hoomd.group.all();
    hoomd.md.integrate.mode_standard(dt=dt);
    langevin = hoomd.md.integrate.langevin(group=all, kT=0.1, seed=987);
    snapshot = system.take_snapshot(all=True)
    print(fpath)
    Eee = []
    with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
        n_frames = f_gsd.nframes
        n_p = f_gsd.read_chunk(frame=0, name='particles/N')
        box = f_gsd.read_chunk(frame=0, name='configuration/box')
        table = hoomd.md.pair.table2D(width, height, 0.5*box[0], 0.5*box[1])
        table.set_from_file(table_path)
        pos = np.zeros((n_frames, n_p[0]))
        pos_frame = f_gsd.read_chunk(frame=0, name='particles/position')
        for j_frame in range(0, n_frames, step):
            pos_frame = f_gsd.read_chunk(frame=j_frame, name='particles/position')

            snapshot.particles.position[:] = pos_frame[:]
            system.restore_snapshot(snapshot)
            hoomd.run(1, quiet=True)
            Eee.append(table.get_energy(hoomd.group.all()))
    return np.array(Eee)

def Eee_Ep_from_gsd(fpath, table_path, width, height, p=None, A=None, phi=None, step = 1):
    """
    Return interaction energy and external periodic potential energy along a trajectory in gsd file located at fpath;
    table - interaction pair force object
    """
    import hoomd
    import hoomd.md
    hoomd.context.initialize('--mode=cpu')
    #hoomd = imp.reload(hoomd)
    system = hoomd.init.read_gsd(fpath, frame = 0)
    if phi==None:
        phi=np.pi
    if p!=None and A!=None:
        periodic = hoomd.md.external.periodic_cos()
        periodic.force_coeff.set('A', A=A, i=0, p=p, phi=phi)
    dt = 0.001
    all = hoomd.group.all();
    hoomd.md.integrate.mode_standard(dt=dt);
    langevin = hoomd.md.integrate.langevin(group=all, kT=0.01, seed=987);
    snapshot = system.take_snapshot(all=True)
    print(fpath)
    Eee = []
    Ep = []
    with gsd.fl.GSDFile(fpath, 'rb') as f_gsd:
        n_frames = f_gsd.nframes
        n_p = f_gsd.read_chunk(frame=0, name='particles/N')
        box = f_gsd.read_chunk(frame=0, name='configuration/box')
        table = hoomd.md.pair.table2D(width, height, 0.5*box[0], 0.5*box[1])
        table.set_from_file(table_path)
        pos = np.zeros((n_frames, n_p[0]))
        pos_frame = f_gsd.read_chunk(frame=0, name='particles/position')
        for j_frame in range(0, n_frames, step):
            pos_frame = f_gsd.read_chunk(frame=j_frame, name='particles/position')
            snapshot.particles.position[:] = pos_frame[:]
            system.restore_snapshot(snapshot)
            hoomd.run(1, quiet=True)
            Eee.append(table.get_energy(hoomd.group.all()))
            Ep.append(periodic.get_energy(hoomd.group.all()))
    return np.array(Eee), np.array(Ep)

def sweep_density(pos, box, window, pts, axis=0):
    """ Density of particles along a chosen axis, computed using sweeping window average.
    \param pos - Nx3 array of particle coordinates
    \param box - hoomd box object with fields Lx and Ly
    \param window - averaging window width
    \param pts - number of points to compute density at
    \param axis - axis (0 or 1) for computing the density along
    """
    if axis == 0:
        L = box.Lx
        h = box.Ly
    elif axis == 1:
        L = box.Ly
        h = box.Lx
    else:
        raise ValueError('axis must be 0 or 1')
    X = np.linspace(-L/2, L/2, pts)
    Y = np.zeros(pts)
    for i, x in enumerate(X):
        Y[i] = len(np.where((pos[:,axis] > x - window/2)*((pos[:,axis] < x + window/2)))[0])
    Y /= (window*h)
    return X, Y

def rotate_positions(pos, box, alpha):
    R = np.array([[np.cos(alpha), -np.sin(alpha), 0], [np.sin(alpha), np.cos(alpha), 0],\
                 [0, 0, 1]])
    rot_pos = np.transpose(np.dot(R, np.transpose(pos)))
    rot_pos = reshape_to_box(rot_pos, np.zeros(3), box)
    return rot_pos

max_stars_repo_name: kmoskovtsev/Electrons-on-Helium-Scripts

max_issues_repo_name: kmoskovtsev/Electrons-on-Helium-Scripts

max_forks_repo_name: kmoskovtsev/Electrons-on-Helium-Scripts

reason: import numpy,from scipy


text:
import math

import dash
import dash_table
from dash import dcc, html
from dash.dependencies import (ALL, MATCH, ClientsideFunction, Input, Output,
                               State)
from dash.exceptions import PreventUpdate
from numpy.lib.function_base import append
from pandas.core.accessor import PandasDelegate

from bsbetl import app_helpers, g
from bsbetl.alltable_calcs import at_columns
from bsbetl.app import app
from bsbetl.app_helpers import at_data_style, init_page_size_input
from bsbetl.func_helpers import (config_shares_list_2_names_or_numbers,
                                 save_and_reload_page_size,
                                 save_runtime_config)
from bsbetl.plots import plot_dashboard_share_chart

layout_at = [
    html.H3(
        children=[
            html.Span(f'ALL_TABLES for period {g.CONFIG_RUNTIME["process-3-last-start"]} to {g.CONFIG_RUNTIME["process-3-last-end"]}', className='border-highlight')],
        className='tab-header'
    ),
    html.Div(
        id= 'dummy',
        children=''
    ),
    html.Div(
        id='conds-tab-summary',
        children='Inspect All-Table calculation results here (post processing). Select a share list and a share... NOTE: Up to 4 share lists (Default, 3jP, V2d, 3nH) should appear once all results processing has also been run',
        className='sw-muted-label'
    ),
    html.Div(children=[
        dcc.Dropdown(
            id='sharelists-dropdown',
            #options=app_helpers.sharelists_options(mastershares_wanted=True),
            options= app_helpers.alltables_sharelists_options(),
            value='Default.shl', # g.CONFIG_RUNTIME['last_processed_sharelist']+'.shl',
            placeholder="Select a sharelist",
            className='dropdown',
            style={'float': 'left', 'width': '260px'},
            persistence=True,
            persistence_type='local',
        ),
        dcc.Dropdown(
            id='shares-dropdown',
            # options=app_helpers.shares_dropdown_options('Default.shl'),
            placeholder="Select a share to view its all-table",
            optionHeight=26,
            value=None,
            className='dropdown',
            style={'float': 'left', 'width': '440px'},
            persistence=True,
            persistence_type='local',
        ),
        html.Label(id='at-share-info-label', children='selection info...', className='share-info'),
    ]),
    html.Div(children=[
        dcc.Slider(
            id='stage-at-slider',
            #tooltip={'always_visible': True, 'placement': 'bottom'},
            min=1,
            max=2,
            step=1,
            value=2,
            className='stages-slider',
            marks={
                1: 'Minutely',
                2: 'Daily',
                # 3: '3',
                # 4: '4',
            },
            included=False,
        ),
    ], style={'clear': 'both'}
    ),
    html.Div(children=[
        dcc.Dropdown(
            id='calcs-at-dropdown',
            options=app_helpers.get_at_calcs_dropdown_options(stage=0),
            placeholder="Column sets per calculation - NOTE calc results appear only in their respective Minutely/Daily All-Table (see slider above). Choose calc(s) , then click 'refresh')",
            multi=True,
            value=None,
            className='dropdown',
            style={'overflow': 'visible', 'width': '100%', },
        ),
    ], style={'clear': 'both'}
    ),
    html.Div(
        className='show-hide-container',
        children=[
            html.Button(
                id='btn-at-fields-reveal',
                className='link-button',
                children='Select|Hide All-Table Columns',
                title='click to reveal|hide column selection checkboxes',
                value='1',
            ),
            html.Button(
                id='at-refresh',
                className='link-button-refresh',
                children='refresh',
                value='1',
                title='click to display new selection of columns',
            ),
            html.Label(id='at-last-date-label', children='',
                       className='last-date-info'),
            dcc.Checklist(
                id='checklist-at-day-ends',
                options=[
                    {'label': 'Day-ends only', 'value': 'dayends'}
                ],
                value=[],
                labelStyle={'display': 'inline'},
                style={'display': 'inline', 'margin-right': '10px'},
            ),
            html.Button(
                id='btn-chart-selected-cols',
                className='link-button-refresh',
                children='Charts for columns selected in the table',
                title='charts those columns whose column headers are selected',
                value='1',
                #style={'float': 'right'}
            ),
            html.Div(id='selections_warning_div', children='Ensure desired column headers are selected',style={'display': 'none', 'margin-left': '10px', 'color': 'firebrick'}),
            html.Button(
                id='btn-lastpage-at',
                children='last pg',
                className='link-button',
                title='go to last page',
                style={'float': 'right'},
            ),
            html.Button(
                id='btn-pagesize-at',
                children='set',
                className='link-button',
                title='set page size',
                style={'float': 'right'},
            ),
            dcc.Input(
                id='at-page-size-inp',
                type='number',
                min=10,
                max=2000,
                step=1,
                value=init_page_size_input('at_page_size'),
                style={'float': 'right'},
            ),
            html.Button(
                id='btn-chart-set',
                className='link-button-refresh',
                children='Charts for selected Calculations',
                title='Charts all selected Calculations',
                value='1',
                style={'float': 'right', 'margin-right': '10px'}
            ),
        ], style={'overflow': 'hidden'}
    ),
    html.Div(
        id='at-containing-div',
        children=[
            html.Div(
                id='status-msg-at',
                className='status-msg',
            ),
            html.Div(
                id='at-fields-div',
                children=[
                    html.Div(children=[
                        dcc.Checklist(
                            id='at-fields-checklist',
                            #options=app_helpers.get_alltable_fields_options(),
                            value=app_helpers.get_alltable_fields_initial_values(),
                            labelStyle={'display': 'inline-block'}
                        ),
                    ], style={'overflow': 'hidden'}
                    ),
                ], style={'display': 'none'} # {'display': 'block', 'overflow': 'hidden'}
            ),
            html.Div(
                id='at-filter-query-div',
                children='',
                className='filter-query',
            ),
            dash_table.DataTable(
                id='share-dash-table',
                # columns=app_helpers.build_share_datatable_columns(),
                #fixed_columns={'headers': True, 'data': 1},
                page_action='native',
                page_current=0,
                page_size=10,
                data=None,
                filter_action='custom',  # 'native'
                filter_query='',
                sort_action='native',
                tooltip_duration=None,
                column_selectable='multi',
                #row_selectable='multi',
                #row_selectable='single',
                selected_columns=[],
                selected_rows=[],
                # style_as_list_view=True,
                style_table={'minWidth': '100%'},
                style_header={
                    'backgroundColor': 'white',
                    'fontWeight': 'bold',
                    'border': '1px solid black',
                },
                style_cell={
                    # all three widths are needed
                    'minWidth': '80px',# 'width': '120px', 'maxWidth': '120px',
                    'overflow': 'hidden',
                    'textOverflow': 'ellipsis',
                },
                style_data_conditional=at_data_style,
                persistence=True,
                persisted_props=['selected_rows'],
                virtualization=True,
            ),
        ], style={'overflow': 'hidden'}
    ),
]

''' callbacks for layout_at '''

@app.callback(
    Output('shares-dropdown', 'options'),
    Input('sharelists-dropdown', 'value'),
    #Input('bourse-radioitems', 'value')]
)
def sharelist_chosen(input_value): #, input1):
    mediated_input = 'Default.shl' if input_value == None else input_value
    bourse_filter_value = g.DEFAULT_BOURSE #if input_value == None else input1
    try:
        ddo = app_helpers.shares_dropdown_options(mediated_input, bourse_filter_value)
        return ddo
    except FileNotFoundError as exc:
        print(f'sharelist {input_value} not found. Have you perhaps started a process-3 again?')
        return dash.no_update

@app.callback(
    Output('at-fields-div', 'style'),
    [Input('btn-at-fields-reveal', 'n_clicks')],
    [State('at-fields-div', 'style')],
    prevent_initial_call=True
)
def at_columns_hide_select_chosen(n_clicks, input1: str):
    if input1 == {'display': 'none'}:
        return {'display': 'block'}
    else:
        return {'display': 'none'}

@app.callback(
    [Output('status-msg-at', 'children'),
     Output('share-dash-table', 'columns'),
     Output('share-dash-table', 'data'),
     Output('share-dash-table', 'tooltip_header'),
     Output('btn-chart-selected-cols', 'disabled'),
     Output('at-last-date-label', 'children'),
     Output('at-filter-query-div', 'children'),
     Output('share-dash-table', 'page_size'),
     Output('calcs-at-dropdown', 'options'),
     #Output('share-dash-table', 'page_current'),
     Output('at-fields-checklist','options'),
     Output('share-dash-table','fixed_columns'),
     Output('share-dash-table','fixed_rows'),
     Output('share-dash-table','row_selectable'),
     Output('at-page-size-inp', 'value')
     ],
    [Input('shares-dropdown', 'value'),
     Input('at-refresh', 'n_clicks'),
     Input('stage-at-slider', 'value'),
     Input('share-dash-table', 'filter_query'),
     ],
    [State('at-fields-checklist', 'value'),
     State('checklist-at-day-ends', 'value'),
     #State('virtual-cols-at-dropdown', 'value'),
     ],
)
def share_chosen(share_name_num: str, n_clicks: int, stage: int, filter,
                 checked_cols: list, dayends_only: bool ): # virtual_cols: list,
    """ extract dataframe from HDFStore
        input_value is string containing share_name followed by share_number eg:
        ANGLO AMERICAN SP.ADR 1/2      A143BR.ETR
        or
        Use ALL Shares                 IGNORE
    """

    #we get our page size not from the input field, but rather from the runtime config
    #and in fact then update the input field with this value to reflect this value
    if stage==1:
        page_size_wanted=516
    else:
        page_size_wanted=g.CONFIG_RUNTIME['at_page_size']

    dash_page_size = page_size_wanted

    fixed_columns={} # { 'headers': True, 'data': 1 }
    fixed_rows={'headers':True, 'data':0}

    fields_chklist_options=app_helpers.get_alltable_fields_options(stage)

    calc_options = app_helpers.get_at_calcs_dropdown_options(stage)

    checked_col_specs = app_helpers.build_share_datatable_columns(checked_cols, stage, None) # virtual_cols

    col_tooltips = app_helpers.build_share_datatable_col_tooltips(checked_cols)

    if isinstance(share_name_num, str) and not share_name_num.endswith('IGNORE'):

        # retrieve and build data structure for the dash datatable - store globally
        at_columns.share_dict = app_helpers.get_share_dataframe_as_dict(
            share_name_num, checked_cols,
            stage=stage, dayends_only=dayends_only, filter=filter, virtual_cols=None)

        if isinstance(at_columns.share_dict, list):

            # date_time assumed ALWAYS selected
            first_dt = at_columns.share_dict[0]['date_time'][:11]
            last_dt = at_columns.share_dict[-1]['date_time'][:11]
            status_msg = f"{first_dt} → {last_dt} (last {len(at_columns.share_dict)} rows ... calculated from a 'start_date' of {g.CONFIG_RUNTIME['last_calculations_start_date']})"
            last_date_label_msg = f'last processed date: {last_dt}'
            # col_options_for_stage,
            return [status_msg, checked_col_specs, at_columns.share_dict, col_tooltips,
                    False, last_date_label_msg, filter, dash_page_size,
                    calc_options, fields_chklist_options, fixed_columns,
                    fixed_rows, 'multi', page_size_wanted]
        else:
            status_msg = f"ERROR: '{share_name_num}' not found in SHARES store, stage {stage}. Inspect '{g.CSV_BY_SHARE_FOLDER}' folder & file of the share for clues. Perhaps no data for the period?"
            # col_options_for_stage,
            #return dash.no_update
            return [status_msg, checked_col_specs, None, None, True,
                    filter, '',  dash_page_size, calc_options,
                    fields_chklist_options, fixed_columns,
                    fixed_rows,'multi', page_size_wanted]
    else:
        # col_options_for_stage,
        return ['Select the AT fields (aphabetically sorted) you would like displayed...',
                None, None, None, True, filter, '',  dash_page_size, calc_options,
                fields_chklist_options, fixed_columns,
                fixed_rows, 'multi', page_size_wanted]

@app.callback(
    Output('at-fields-checklist', 'value'),
    Input('calcs-at-dropdown', 'value')
)
def at_calculations_chosen(calcs_selected):

    #print(f'calcs_selected={calcs_selected}')
    if isinstance(calcs_selected, list) and len(calcs_selected) > 0:
        wanted_cols = ['date_time', 'price', 'volume']
        for cols_for_calc in calcs_selected:
            for col in cols_for_calc.split(','):
                wanted_cols.append(col)

        # print(wanted_cols)
        return wanted_cols
    else:
        # init_cols = app_helpers.get_alltable_fields_initial_values()
        # return init_cols
        raise PreventUpdate

@app.callback(
    Output('local-store', 'data'),
    Input({'type': 'input-parameter', 'index': ALL}, 'value'),
    [State({'type': 'input-parameter', 'index': ALL}, 'id'),
     State('local-store', 'data')
     ],
    prevent_initial_call=True
)
def filter_parameter_changed(value, id, data):
    ''' fires when plugin filter parameter is changed by the user
    NOTE we can get here from EITHER the all-tables OR overview tabs
    '''

    # print(f'filter_parameter_changed in. id={id}, data={data}')
    data = data or {'dummy_key': 0}
    for i, param_dict in enumerate(id):
        # data[id[0]['index']] = value[0]
        data[param_dict['index']] = value[i]

    #print(f'filter_parameter_changed out data={data}')
    return data

@app.callback(
    Output('selections_warning_div','style'),
    #Output("hidden_div_for_redirect_callback", "children"),
    Input('btn-chart-selected-cols', 'n_clicks'),
    [State('shares-dropdown', 'value'),
     State('share-dash-table', 'selected_columns'),
     State('stage-at-slider', 'value')
     ],
    prevent_initial_call=True
)
def selected_columns_chart_chosen(input_value: str, share_name_num: str, cols_to_plot: list, stage: int):
    ''' create a plot and get it displayed '''

    # (regular) cols_to_plot must be mediated by stage
    stage_cols = [
        col for col in cols_to_plot if (col in at_columns.AT_STAGE_TO_COLS[stage])] # or (col in virt_cols)]

    if len(cols_to_plot)>0 and isinstance(share_name_num, str) and len(share_name_num) > 0:
        # print(stage_cols)
        # print(at_columns.AT_COLS_AFFINITY)
        last_plot_tuple = plot_dashboard_share_chart(share_name_num, stage_cols, stage, data_source='local-store', use_config_date_range=True)
        app_helpers.assign_global_at_plot_values(last_plot_tuple)
        #print(f'selected_columns_chart_chosen, type(g.at_plot_df_datetime_col)={type(g.at_plot_df_datetime_col)}')
        #print(f'selected_columns_chart_chosen={g.at_plot_df_datetime_col}')
        return {'display': 'none', 'margin-left': '10px', 'color': 'firebrick'}
        #return dcc.Location(pathname='/apps/at_charts',id='does not matter')
    else:
        return {'display': 'inline', 'margin-left': '10px', 'color': 'firebrick'}
        #raise PreventUpdate

@app.callback(
    Output('dummy', 'children'),
    [Input('btn-chart-set', 'n_clicks'),
     State('calcs-at-dropdown', 'value')
    ],
    [State('shares-dropdown', 'value'),
     #State('share-dash-table', 'selected_columns'),
     State('stage-at-slider', 'value'),
     #State('virtual-cols-at-dropdown', 'value'),
     ],
    prevent_initial_call=True
)
def selected_calculations_chart_chosen(input_value: str, calcs_selected, share_name_num: str, stage: int):
    ''' create a plot and get it displayed '''

    #print(f'calcs_selected={calcs_selected}')
    cols_to_plot = ['date_time', 'price', 'volume']
    if isinstance(calcs_selected, list) and len(calcs_selected) > 0:
        for cols_for_calc in calcs_selected:
            for col in cols_for_calc.split(','):
                cols_to_plot.append(col)
    else:
        return ''

    # (regular) cols_to_plot must be mediated by stage
    stage_cols = [
        col for col in cols_to_plot if (col in at_columns.AT_STAGE_TO_COLS[stage])] # or (col in virt_cols)]

    #print(f'stage_cols={stage_cols}')

    if isinstance(share_name_num, str) and len(share_name_num) > 0:
        #print(f'stage_cols={stage_cols}')
        # print(at_columns.AT_COLS_AFFINITY)
        last_plot_tuple = plot_dashboard_share_chart(share_name_num, stage_cols, stage, data_source='hdf', use_config_date_range=True) #'hdf'
        app_helpers.assign_global_at_plot_values(last_plot_tuple)
        return ''
    else:
        return ''

@app.callback(
    Output('at-refresh','n_clicks'),
    Input('btn-pagesize-at','n_clicks'),
    State('at-page-size-inp','value'),
    prevent_initial_call=True
)
def set_at_page_size(n_clicks,page_size):

    save_and_reload_page_size('at_page_size', page_size)
    # this should cause a refresh
    return 0
    #return dash.no_update

@app.callback(
    Output('share-dash-table','page_current'),
    Input('btn-lastpage-at','n_clicks'),
    State('share-dash-table','data'),
    State('at-page-size-inp','value'),
)
def go_to_last_page(n_clicks, data, cur_page_size):

    if data is not None:
        last_page = math.ceil(len(data) / cur_page_size) - 1
        #print(last_page)
        return last_page
    else:
        return dash.no_update

@app.callback(
    Output('share-dash-table', "style_data_conditional"),
    Input('share-dash-table', "selected_rows")
)
def style_selected_rows(sel_rows):

    at_hilite_style = [ {"if": {"filter_query": "{{id}}={}".format(i)}, "backgroundColor": "lightgoldenrodyellow",} for i in sel_rows ]

    if sel_rows is None or len(sel_rows)==0:
        return dash.no_update

    val = at_data_style + at_hilite_style

    return val

@app.callback(
    Output('at-share-info-label','children'),
    Input('shares-dropdown', 'value'),
    prevent_initial_call = True
)
def update_share_info(share_and_number):

    sharelists_list = []
    _3jP_list = config_shares_list_2_names_or_numbers(g.CONFIG_RUNTIME['_3jP_list'],numbers_wanted=True)
    _V2d_list = config_shares_list_2_names_or_numbers(g.CONFIG_RUNTIME['_V2d_list'],numbers_wanted=True)
    _3nH_list = config_shares_list_2_names_or_numbers(g.CONFIG_RUNTIME['_3nH_list'],numbers_wanted=True)

    #Aareal Bank AG                 540811.ETR
    share_num = share_and_number[-11:-1] # there's a \n at the end of the line

    if share_num in _3jP_list:
        sharelists_list.append('3jP')
    if share_num in _V2d_list:
        sharelists_list.append('V2d')
    if share_num in _3nH_list:
        sharelists_list.append('3nH')

    inform=''
    if len(sharelists_list)>0:
        inform = f"this share is currently also in the following 'selection sharelist(s)': {', '.join(sharelists_list)}"

    return inform

reason: from numpy


text:
"""Unit test package for backtest."""
import logging
import os
import pickle
import socket
import uuid
from contextlib import closing
from typing import Union

import arrow
import cfg4py
import numpy as np
import pandas as pd
import pkg_resources
from coretypes import FrameType, bars_dtype
from omicron.dal.influx.influxclient import InfluxClient
from omicron.models.stock import Stock
from omicron.models.timeframe import TimeFrame

from backtest.app import application as app
from backtest.common.errors import BacktestError
from backtest.config import endpoint, get_config_dir
from backtest.web.interfaces import bp

os.environ[cfg4py.envar] = "DEV"
cfg = cfg4py.init(get_config_dir())
logger = logging.getLogger(__name__)

def init_interface_test():
    bp.url_prefix = endpoint()
    app.blueprint(bp)

    return app

def find_free_port():
    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
        s.bind(("localhost", 0))
        # s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        return s.getsockname()[1]

async def delete(cmd: str, token: str, params=None):
    url = f"{endpoint()}/{cmd}"

    headers = {"Authorization": f"Token {token}", "Request-ID": uuid.uuid4().hex}
    _, response = await app.asgi_client.delete(url, params=params, headers=headers)
    content_type = response.headers.get("Content-Type")

    if response.status == 200:
        if content_type == "application/json":
            return response.json
        elif content_type.startswith("text"):
            return response.text
        else:
            return pickle.loads(response.content)
    if response.status == 499:
        raise BacktestError(response.status, response.text)

async def post(cmd: str, token: str, data):
    url = f"{endpoint()}/{cmd}"

    headers = {"Authorization": f"Token {token}", "Request-ID": uuid.uuid4().hex}
    _, response = await app.asgi_client.post(url, json=data, headers=headers)
    content_type = response.headers.get("Content-Type")

    if response.status == 200:
        if content_type == "application/json":
            return response.json
        elif content_type.startswith("text"):
            return response.text
        else:
            return pickle.loads(response.content)
    if response.status == 499:
        raise BacktestError(response.status, response.text)

async def get(cmd: str, token: str, **kwargs):
    url = f"{endpoint()}/{cmd}"

    headers = {"Authorization": f"Token {token}", "Request-ID": uuid.uuid4().hex}
    _, response = await app.asgi_client.get(url, headers=headers, params=kwargs)

    content_type = response.headers.get("Content-Type")

    if response.status == 200:
        if content_type == "application/json":
            return response.json
        elif content_type.startswith("text"):
            return response.text
        else:
            return pickle.loads(response.content)
    if response.status == 499:
        raise BacktestError(response.status, response.text)

def data_dir():
    return os.path.join(os.path.dirname(__file__), "data")

def read_csv(fname, start=None, end=None):
    """start, end是行计数，从1开始，以便于与编辑器展示的相一致。
    返回[start, end]之间的行
    """
    path = os.path.join(data_dir(), fname)
    with open(path, "r") as f:
        lines = f.readlines()

    if start is None:
        start = 1  # skip header
    else:
        start -= 1

    if end is None:
        end = len(lines)

    return lines[start:end]

def lines2bars(lines, is_date):
    """将CSV记录转换为Bar对象

    header: date,open,high,low,close,money,volume,factor
    lines: 2022-02-10 10:06:00,16.87,16.89,16.87,16.88,4105065.000000,243200.000000,121.719130

    """
    if isinstance(lines, str):
        lines = [lines]

    def parse_date(x):
        return arrow.get(x).date()

    def parse_naive(x):
        return arrow.get(x).naive

    if is_date:
        convert = parse_date
    else:
        convert = parse_naive

    data = []
    for line in lines:
        fields = line.split(",")
        data.append(
            (
                convert(fields[0]),
                float(fields[1]),
                float(fields[2]),
                float(fields[3]),
                float(fields[4]),
                float(fields[5]),
                float(fields[6]),
                float(fields[7]),
            )
        )

    return np.array(data, dtype=bars_dtype)

def bars_from_csv(
    code: str, ft: Union[str, FrameType], start_line: int = None, end_line: int = None
):
    ft = FrameType(ft)

    fname = f"{code}.{ft.value}.csv"

    if ft in TimeFrame.minute_level_frames:
        is_date = False
    else:
        is_date = True

    return lines2bars(read_csv(fname, start_line, end_line), is_date)

def assert_deep_almost_equal(test_case, expected, actual, *args, **kwargs):
    """
    copied from https://github.com/larsbutler/oq-engine/blob/master/tests/utils/helpers.py

    Assert that two complex structures have almost equal contents.
    Compares lists, dicts and tuples recursively. Checks numeric values
    using test_case's :py:meth:`unittest.TestCase.assertAlmostEqual` and
    checks all other values with :py:meth:`unittest.TestCase.assertEqual`.
    Accepts additional positional and keyword arguments and pass those
    intact to assertAlmostEqual() (that's how you specify comparison
    precision).
    :param test_case: TestCase object on which we can call all of the basic
        'assert' methods.
    :type test_case: :py:class:`unittest.TestCase` object
    """
    is_root = "__trace" not in kwargs
    trace = kwargs.pop("__trace", "ROOT")
    try:
        if isinstance(expected, (int, float, complex)):
            test_case.assertAlmostEqual(expected, actual, *args, **kwargs)
        elif isinstance(expected, (list, tuple, np.ndarray)):
            test_case.assertEqual(len(expected), len(actual))
            for index in range(len(expected)):
                v1, v2 = expected[index], actual[index]
                assert_deep_almost_equal(
                    test_case, v1, v2, __trace=repr(index), *args, **kwargs
                )
        elif isinstance(expected, dict):
            test_case.assertEqual(set(expected), set(actual))
            for key in expected:
                assert_deep_almost_equal(
                    test_case,
                    expected[key],
                    actual[key],
                    __trace=repr(key),
                    *args,
                    **kwargs,
                )
        else:
            test_case.assertEqual(expected, actual)
    except AssertionError as exc:
        exc.__dict__.setdefault("traces", []).append(trace)
        if is_root:
            trace = " -> ".join(reversed(exc.traces))
            exc = AssertionError("%s\nTRACE: %s" % (exc, trace))
        raise exc

async def data_populate():
    cfg = cfg4py.init(get_config_dir())
    url, token, bucket, org = (
        cfg.influxdb.url,
        cfg.influxdb.token,
        cfg.influxdb.bucket_name,
        cfg.influxdb.org,
    )
    client = InfluxClient(url, token, bucket, org)

    # fill in influxdb
    await client.drop_measurement("stock_bars_1d")
    await client.drop_measurement("stock_bars_1m")

    for ft in (FrameType.MIN1, FrameType.DAY):
        file = os.path.join(data_dir(), f"bars_{ft.value}.pkl")
        with open(file, "rb") as f:
            bars = pickle.load(f)
            await Stock.persist_bars(ft, bars)

    df = pd.read_csv(
        os.path.join(data_dir(), "limits.csv"), sep="\t", parse_dates=["time"]
    )
    limits = df.to_records(index=False)
    limits.dtype.names = ["frame", "code", "high_limit", "low_limit"]
    await Stock.save_trade_price_limits(limits, False)

def _index_of(array, date):
    idx = np.argwhere(array["date"] <= date).ravel()
    assert len(idx) <= 1
    if len(idx) == 0:
        return None
    else:
        return idx[0]

max_stars_repo_name: zillionare/backtesting

max_issues_repo_name: zillionare/backtesting

max_forks_repo_name: zillionare/backtesting


text:
from pyart.core.transforms import cartesian_to_geographic
import xarray as xr
import numpy as np
import pandas as pd
import calendar
from tint.grid_utils import parse_grid_datetime

def get_datetime_components(date_time):
    dt_index = pd.DatetimeIndex([date_time])
    components = []
    for block in ['year', 'month', 'day', 'hour', 'minute']:
        components.append(
            eval('dt_index.{}.values[0]'.format(block)))
    return components

def get_ERA5_daterange(year, month):
    last_day = calendar.monthrange(year, month)[1]
    date_range = '{}{}{}-{}{}{}'.format(
        year, str(month).zfill(2), '01',
        year, str(month).zfill(2), last_day)
    return date_range

def collect_ERA5_files(base_dir, year, range_str, files=[]):
    for variable in ['z', 'u', 'v']:
        fn = '{}/{}/{}/{}_era5_oper_pl_{}.nc'.format(
            base_dir, variable, year, variable, range_str)
        files.append(fn)
    return files

def get_ERA5_ds(
        start_datetime, end_datetime,
        time_delta=np.timedelta64(10, 'm'),
        base_dir='/g/data/rt52/era5/pressure-levels/reanalysis/'):
    s_comps = get_datetime_components(start_datetime)
    [s_year, s_month, s_day, s_hour, s_minute] = s_comps
    e_comps = get_datetime_components(end_datetime)
    [e_year, e_month, e_day, e_hour, e_minute] = e_comps

    if not (e_day == 1 and e_hour == 0 and e_minute == 0):
        end_datetime = end_datetime + time_delta
        e_comps = get_datetime_components(end_datetime)
        [e_year, e_month, e_day, e_hour, e_minute] = e_comps

    start_month = np.datetime64('{}-{}'.format(
        str(s_year).zfill(2), str(s_month).zfill(2)))
    end_month = np.datetime64('{}-{}'.format(
        str(e_year).zfill(2), str(e_month).zfill(2)))
    month_range = np.arange(
        start_month, end_month+np.timedelta64(1, 'M'),
        np.timedelta64(1, 'M'))

    files = []
    for month_dt in month_range:
        year, month = get_datetime_components(month_dt)[:2]
        range_str = get_ERA5_daterange(year, month)
        files = collect_ERA5_files(
            base_dir, year, range_str, files=files)

    return xr.open_mfdataset(files)

def flexible_round(x, prec=2, base=.05, method=round):
    return round(base * method(float(x) / base), prec)

def interp_ERA_ds(ds_all, grid, timedelta=np.timedelta64(10, 'm')):
    lon, lat = cartesian_to_geographic(
        grid.x['data'].data, grid.y['data'].data, grid.get_projparams())
    alt = grid.z['data'].data

    min_lon = flexible_round(min(lon), prec=1, base=.2, method=np.floor)
    max_lon = flexible_round(max(lon), prec=1, base=.2, method=np.ceil)
    min_lat = flexible_round(min(lat), prec=2, base=.25, method=np.floor)
    max_lat = flexible_round(max(lat), prec=2, base=.25, method=np.ceil)

    grid_time = parse_grid_datetime(grid)
    start_time = np.datetime64(grid_time.replace(minute=0, second=0))
    end_time = start_time + np.timedelta64(1, 'h')

    ds = ds_all.loc[dict(
        latitude=slice(max_lat+.25, min_lat-.25),
        longitude=slice(min_lon-.2, max_lon+.2),
        time=slice(start_time, end_time))]
    ds['z'] = ds['z'] / 9.80665
    altitude = ds['z'].mean(['longitude', 'latitude', 'time'])
    ds = ds.assign_coords({'level': altitude})
    ds = ds.rename({'level': 'altitude'})
    ds = ds.drop_vars('z')
    ds = ds.loc[dict(altitude=slice(22000, 0))]
    times = np.arange(start_time, end_time, timedelta)
    ds = ds.interp(longitude=lon, latitude=lat, altitude=alt, time=times)
    return ds

def init_ERA5(grid, params):
    grid_datetime = parse_grid_datetime(grid)
    grid_datetime = np.datetime64(grid_datetime.replace(second=0))
    print('Getting ERA5 metadata.')
    ERA5_all = get_ERA5_ds(
        grid_datetime, grid_datetime+np.timedelta64(1, 'h'),
        base_dir=params['AMBIENT_BASE_DIR'])
    print('Getting Intepolated ERA5 for next hour.')
    ERA5_interp = interp_ERA_ds(
        ERA5_all, grid, timedelta=np.timedelta64(10, 'm'))
    ERA5_interp.load()
    return ERA5_all, ERA5_interp

def update_ERA5(grid, params, ERA5_all, ERA5_interp):
    grid_datetime = parse_grid_datetime(grid)
    grid_datetime = np.datetime64(grid_datetime.replace(second=0))
    grid_datetime_hour = parse_grid_datetime(grid)
    grid_datetime_hour = np.datetime64(
        grid_datetime_hour.replace(second=0, minute=0))
    t_min = ERA5_all.time.values.min()
    t_max = ERA5_all.time.values.max()
    if not (
            grid_datetime >= t_min
            and grid_datetime_hour + np.timedelta64(1, 'h') <= t_max):
        print('Getting ERA5 metadata.')
        ERA5_all = get_ERA5_ds(
            grid_datetime_hour, grid_datetime_hour+np.timedelta64(1, 'h'),
            base_dir=params['AMBIENT_BASE_DIR'])
    if grid_datetime not in ERA5_interp.time.values:
        print('Getting Intepolated ERA5 for the next hour.')
        ERA5_interp = interp_ERA_ds(
            ERA5_all, grid, timedelta=np.timedelta64(10, 'm'))
        ERA5_interp.load()
    return ERA5_all, ERA5_interp


text:
# ½-day to develop. ¼-day to debug.
import torch
import asyncio
import numpy as np

# Input: gather state along streams, and put observations into it.
def webenv_ignore(state, obs_t):
  """
  Ignores observations. Loss will take them into account anyway.
  """
  return state
def webenv_add(state, obs_t):
  """
  Teacher Forcing.
  Adds WebEnv observations to internal state, properly ignoring holes (NaNs).
  """
  padded = torch.nn.functional.pad(obs_t, (0, state.shape[-1] - obs_t.shape[-1]), value=np.nan)
  return torch.where(torch.isnan(padded), state, state + padded)
def webenv_merge(state, obs_t):
  """
  Merges (writes) WebEnv observations into internal state, properly ignoring holes (NaNs).

  Note that this overwrites predictions, so the model cannot access them. If you want your model to know both real and predicted numbers, use `input=webenv_concat` in `recurrent`.
  """
  padded = torch.nn.functional.pad(obs_t, (0, state.shape[-1] - obs_t.shape[-1], 0,0), value=np.nan)
  return torch.clamp(torch.where(torch.isnan(padded), state, padded), -1, 1)
def webenv_concat(state, obs_t):
  """
  Puts observations after predictions. Expect twice the numbers in your transition model.

  Holes (NaNs) in observations will be replaced with the prediction, allowing observation resizing at the cost of repetition.
  """
  obs2 = webenv_merge(state, obs_t)
  return torch.cat((state, obs2), dim=-1)

# Output.
async def webenv_slice(lock, state, indices, obs, act_len):
  """
  Turns two PyTorch state tensors (pre-step and pre-output) (and what WebEnv received: observations and action lengths) into post-step state and lists of predictions and actions, asynchronously.
  """
  # `obs` and `act_len` are lists of equal length `state.shape[0]`.
  with torch.no_grad():
    max_slice = state.shape[-1] # Only allow up to 100% in a slice.
    preds, acts = [], []
    for i in range(len(obs)):
      # Slice obs and actions from ends. (Reverse actions for stability.)
      ind = indices[i, 0]
      pred_t = state[ind, 0:min(max_slice, obs[i].shape[-1])]
      act_t = state[ind, max_slice - min(max_slice, act_len[i]):max_slice].flip(-1)
      # Asynchronously copy to CPU.
      pred = torch.zeros_like(pred_t, layout=torch.strided, device='cpu', memory_format=torch.contiguous_format)
      act = torch.zeros_like(act_t, layout=torch.strided, device='cpu', memory_format=torch.contiguous_format)
      pred.copy_(pred_t, non_blocking=True)
      act.copy_(act_t, non_blocking=True)
      preds.append(pred)
      acts.append(act)
    lock.set_result(None)
    # Wait until all GPU→CPU copies are done, then return lists of predictions and actions.
    event = torch.cuda.Event()
    event.record()
    while not event.query():
      await asyncio.sleep(.001)
    return [p.numpy() for p in preds], [a.numpy() for a in acts]

# Pre-input and pre-output.
def webenv_gather(state1, indices): # → state2
  if state1.shape[0] == 1: return state1
  indices = torch.tensor(indices, device = state1.device)
  indices = indices.expand(indices.shape[0], *state1.shape[1:]) # Apparently required for the backward pass to work.
  return torch.gather(state1, 0, indices)
def webenv_scatter(state1, indices, state2): # → state1
  if state1.shape[-1] != state2.shape[-1]:
    raise TypeError('Pre/post transition sizes mismatch')
  if state1.shape[0] == 1: return state2
  indices = torch.tensor(indices, device = state1.device) # Double op; inefficient.
  indices = indices.expand(indices.shape[0], *state1.shape[1:]) # Apparently required for the backward pass to work:
  #   https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_.html
  return state1.clone().scatter_(0, indices, state2)

# Minimization targets.
def L1(pred, got, *_):
  return (pred - got).abs().sum()
def L2(pred, got, *_):
  return .5*(pred - got).square().sum()

def optimizer_step(optimizer):
  optimizer.step()
  optimizer.zero_grad()

# RNN.
def recurrent(
  state,
  loss = L2,
  optimizer=None,
  device='cuda',
  unroll_length=16,
  unrolls_per_step=1,
  synth_grad=None,
  synth_grad_loss = L2,
  input = webenv_merge,
  output = webenv_slice,
  gather = webenv_gather,
  scatter = webenv_scatter,
  update = optimizer_step,
):
  """
  Creates a decorator, which creates a real-time recurrent multi-stream transformer (input→output).

  Args:
    `state`: the initial state or its shape, such as `(1,64)`. The decorated `transition` takes a state and returns a state, as PyTorch tensors.
      (This is never re-allocated, so make sure to never go above this.)
    `loss`: computes the number to minimize, given `pred` and `actual` (and all args). L2 by default.
      (`pred` and `actual` differ only in `input`. The shared parts can be conditioned-on, by learned losses.)
    `optimizer`: the PyTorch optimizer. Adam by default.
    `device`: the device to use for numeric computations. `'cuda'` by default.
    `unroll_length`: how long to accumulate gradients before applying them, either a number or a function that takes a number and returns a bool. 16 by default.
      (Backpropagation-through-time.)
      (A high unroll length adds a lot of latency to some frames, because PyTorch has no easy way to desynchronize the backward pass at the cost of some correctness.)
      (If 1, specify `synth_grad`.)
    `unroll_per_step`: how many unrolls per `optimizer` step. 1 by default.
      (This is technically the upper bound, because not all streams have data available at every step.)
    `synth_grad`: a function from state to its gradient, which stitches different unrolls together. `None` by default.
    `synth_grad_loss`: computes the number to minimize. L2 by default.
    `input`: goes from PyTorch state and observation to the updated state. `webenv_merge` by default.
    `output`: goes from PyTorch state and step's args to the output, async. `webenv_slice` by default.
    `gather`: extracts stream state slices before `input`. `webenv_gather` by default.
    `scatter`: reunites stream state slices after `output`. `webenv_scatter` by default.
    `update`: applies parameter updates, given the optimizer. `optimizer_step` by default.
  After these args, supply the `transition` in another call.
  Then, await calls to step, passing in indices (`0` to only have one stream, else `np.array([[0],[1],[3],[4]], dtype=np.int64)`), observations (NaN-filled where lengths mismatch), and any other args.
  """
  def rec(transition):
    nonlocal optimizer, state
    if optimizer is None:
      optimizer = torch.optim.Adam(transition.parameters(), lr=3e-4)
    if isinstance(state, list) or isinstance(state, tuple):
      if len(state) != 2:
        raise TypeError('State must be 2D')
      state = torch.zeros(*state, device=device)
    else:
      state = state.copy(device=device)
    state.detach_().requires_grad_()
    start_state = state
    unroll_index = 0
    unroll_loss = 0
    unrolls = 0
    async def step(lock, indices, obs, *args):
      # Step.
      nonlocal start_state, state, unroll_index, unroll_loss, unrolls
      if indices.max() >= state.shape[0]:
        raise TypeError('Got too many streams: got ' + str(indices.max()+1) + ' but only have state for ' + str(state.shape[0]))
      obs_t = list_to_torch(obs, device)
      state2 = gather(state, indices)
      state3 = input(state2, obs_t)
      # Prev frame predicts this one:
      unroll_loss = unroll_loss + loss(state2, (webenv_merge(state2, obs_t) if input is not webenv_merge else state3).detach(), obs_t, *args)

      state4 = transition(state3)
      state = scatter(state, indices, state4)
      unroll_index += 1
      # Backprop.
      if unroll_length(unroll_index) if callable(unroll_length) else (unroll_length <= unroll_index):
        if synth_grad is not None:
          with torch.no_grad():
            grad = state - synth_grad(state)
          unroll_loss = unroll_loss + (state * grad).sum()
        unroll_loss.backward()
        unroll_loss = 0.
        if synth_grad is not None:
          st = start_state.detach()
          synth_grad_loss(synth_grad(st), st - start_state.grad).backward()
        unrolls += 1
        if unrolls >= unrolls_per_step:
          unrolls = 0
          update(optimizer)
          optimizer.step()
          optimizer.zero_grad()
        start_state = state = state.detach().requires_grad_(True)
        unroll_index = 0
      return await output(lock, state, indices, obs, *args)
    return step
  return rec
def list_to_torch(xs, device):
  # NaN-pad these 1D NumPy arrays to their max length and stack, then send to `device`.
  max_dim = max(x.shape[-1] for x in xs)
  xs = np.stack([np.pad(x, (0, max_dim - x.shape[-1]), constant_values=np.nan) for x in xs])
  return torch.tensor(xs, device=device)

if __name__ == '__main__':
  async def test(n, device):
    obs = np.random.randn(1, 32).astype('float32') * .1
    obs[0,2] = np.nan

    class Dense(torch.nn.Module):
      def __init__(self, N, device):
        super(Dense, self).__init__()
        p = torch.nn.Parameter
        self.weights = p(torch.randn(N, N, device=device), requires_grad=True)
        self.bias = p(torch.randn(N, device=device), requires_grad=True)
        self.scale = p(torch.randn(N, device=device), requires_grad=True)
      def forward(self, x):
        x = torch.matmul(x, self.weights.data.requires_grad_(True))
        x = x / (x.square().sum().sqrt() + 1e-12)
        x = x * self.scale
        x = x + self.bias
        return x

    N = 128
    transition = Dense(N, device)
    synth = Dense(N, device)
    optim = torch.optim.Adam([*transition.parameters(), *synth.parameters()], lr=.001)
    stream = recurrent(
      (1,N), device=device, optimizer=optim, unroll_length=2, synth_grad=synth,
    )(transition)

    import time
    start = time.time()
    indices = np.array([[0]], dtype=np.int64)
    for i in range(n):
      preds, acts = await stream(asyncio.Future(), indices, obs, [0])
      with torch.no_grad():
        print('L1:', np.nansum(np.abs(preds[0] - obs)))
    print('Time:', time.time() - start, 's')
  asyncio.run(test(5000, 'cuda'))

max_stars_repo_name: Antipurity/webenv

max_issues_repo_name: Antipurity/webenv

max_forks_repo_name: Antipurity/webenv