text:
from pathlib import Path
import numpy as np

from ..units import AMBER_HARTREE_TO_KCAL, AMBER_BOHR_TO_A
from .templates.sqm import get_qm_template, default_options
from .qmbase import QMBase

class SQM(QMBase):

    OUTPUT = "sqm.out"
    default_options = default_options

    def gen_input(self):
        """Generate input file for QM software."""

        if not "qmcharge" in self.options:
            self.options["qmcharge"] = str(self.charge)

        if not "spin" in self.options:
            self.options["spin"] = str(self.mult)

        with open(Path(self.cwd).joinpath("sqm.inp"), 'w') as f:
            f.write(get_qm_template(self.options))

            for e, s, x, y, z, in zip(
                self.qm_elements,
                self.qm_element_symbols,
                self.qm_positions[0],
                self.qm_positions[1],
                self.qm_positions[2],
            ):
                f.write(f"{e:3} {s:>2} {x:21.14e} {y:21.14e} {z:21.14e}\n")

            if self.mm_charges is not None:
                f.write("#EXCHARGES\n")
                for x, y, z, c in zip(
                    self.mm_positions[0],
                    self.mm_positions[1],
                    self.mm_positions[2],
                    self.mm_charges,
                ):
                    f.write(f"  1  H {x:21.14e} {y:21.14e} {z:21.14e} {c:21.14e}\n")
                f.write("#END\n")

    def gen_cmdline(self):
        """Generate commandline for QM calculation."""

        cmdline = f"cd {self.cwd}; "
        cmdline += "sqm -O -i sqm.inp -o sqm.out"

        return cmdline

    def _get_qm_energy(self, qm_cache=None, output=None):
        """Get QM energy from output of QM calculation."""

        if qm_cache is not None:
            output = qm_cache
        else:
            try:
                output = Path(output).read_text().split("\n")
            except:
                output = Path(self.cwd).joinpath(self.OUTPUT).read_text().split("\n")
            else:
                raise ValueError("Can not open output.")

        for line in output:
            if "QMMM: SCF Energy" in line:
                return float(line.split()[4]) / AMBER_HARTREE_TO_KCAL

    def _get_qm_energy_gradient(self, qm_cache=None, output=None):
        """Get QM energy gradient from output of QM calculation."""

        if qm_cache is not None:
            output = qm_cache
        else:
            try:
                output = Path(output).read_text().split("\n")
            except:
                output = Path(self.cwd).joinpath(self.OUTPUT).read_text().split("\n")
            else:
                raise ValueError("Can not open output.")

        for i in range(len(output)):
            if "Forces on QM atoms from SCF calculation" in output[i]:
                gradient = np.empty((3, len(self.qm_elements)), dtype=float)
                for j in range(len(self.qm_elements)):
                    line = output[i + j + 1]
                    gradient[0, j] = float(line[18:38])
                    gradient[1, j] = float(line[38:58])
                    gradient[2, j] = float(line[58:78])
                return gradient / (AMBER_HARTREE_TO_KCAL / AMBER_BOHR_TO_A)

    def _get_mm_esp(self, qm_cache=None, output=None):
        """Get electrostatic potential at MM atoms in the near field from QM density."""

        if qm_cache is not None:
            output = qm_cache
        else:
            try:
                output = Path(output).read_text().split("\n")
            except:
                output = Path(self.cwd).joinpath(self.OUTPUT).read_text().split("\n")
            else:
                raise ValueError("Can not open output.")

        mm_esp = np.zeros((4, len(self.mm_charges)))

        for i in range(len(output)):
            if "Forces on MM atoms from SCF calculation" in output[i]:
                for j in range(len(self.mm_charges)):
                    line = output[i + j + 1]
                    mm_esp[1, j] = float(line[18:38]) / self.mm_charges[j]
                    mm_esp[2, j] = float(line[38:58]) / self.mm_charges[j]
                    mm_esp[3, j] = float(line[58:78]) / self.mm_charges[j]
                line_num = i + j + 2
                break

        for i in range(line_num, len(output)):
            if "Electrostatic Potential on MM atoms from QM Atoms" in output[i]:
                for j in range(len(self.mm_charges)):
                    line = output[i + j + 1]
                    mm_esp[0, j] = float(line[18:38])
                break

        mm_esp[0] /= AMBER_HARTREE_TO_KCAL
        mm_esp[1:] /= (AMBER_HARTREE_TO_KCAL / AMBER_BOHR_TO_A)
        return mm_esp

    def _get_mulliken_charges(self, qm_cache=None, output=None):
        """Get Mulliken charges from output of QM calculation."""

        if qm_cache is not None:
            output = qm_cache
        else:
            try:
                output = Path(output).read_text().split("\n")
            except:
                output = Path(self.cwd).joinpath(self.OUTPUT).read_text().split("\n")
            else:
                raise ValueError("Can not open output.")

        for i in range(len(output)):
            if "Atomic Charges" in output[i]:
                mulliken_charges = np.empty(len(self.qm_elements), dtype=float)
                for j in range(len(self.qm_elements)):
                    line = output[i + j + 2]
                    mulliken_charges[j] = float(line.split()[-1])
                return mulliken_charges

max_stars_repo_name: panxl/qmhub

max_issues_repo_name: panxl/qmhub

max_forks_repo_name: panxl/qmhub


text:
from pydub import AudioSegment
import numpy, scipy, matplotlib.pyplot as plt, sklearn, librosa, mir_eval, urllib
import os

def splice(audioFile, times):
	audio = AudioSegment.from_wav(audioFile)
	syllables = []
	for i in range(len(times)):
		t = times[i]
		t = t * 1000 #set to milliseconds, might have to drop sig figs
		if i == 0:
			syllables.append(audio[:t])
		else:
			t_prev = times[i - 1] * 1000
			syllables.append(audio[t_prev:t])
			print str(t_prev) + " " + str(t)

		if i == len(times) - 1:
			syllables.append(audio[t:])
	return syllables

def get_times(timeFile):
	times = []
	with open(timeFile) as f:
		for line in f:
			times.append(float(line.strip()))
	return times

def splice_audio(inputDir, inputPathName, onset_times):

	index = 0
	raw_syllables = list()
	#print times
	for filename in os.listdir(inputDir + "/" + inputPathName + '/'):
	    if filename.endswith(".wav"):

		syllables = splice(inputDir + '/' + inputPathName + '/' + filename, onset_times[index])

		for i, syllable in enumerate (syllables):
			out_file = "./audio/output/syllables/" + inputPathName + "/" + filename.rstrip('.wav') + "{0}.wav".format(i)
			if not os.path.exists(os.path.dirname(out_file)):
			    try:
			        os.makedirs(os.path.dirname(out_file))
			    except OSError as exc: # Guard against race condition
			        if exc.errno != errno.EEXIST:
			            raise
			syllable.export(out_file, format="wav")
			x, fs = librosa.load(out_file)
			raw_syllables.append(x)

		index += 1

	#return syllables as vector
	return raw_syllables

ext: py

reason: import numpy


text:
# %%
from init_objects import *
from mpl_toolkits.mplot3d.axes3d import Axes3D
import pandas as pd
import subprocess
import matplotlib.animation as animation
import mpl_toolkits.mplot3d.axes3d as p3
import scipy.integrate
from matplotlib import cm
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d
import numpy as np
from Objects import Objects
import copy
import time
from solar import SolarSystem
from scipy.integrate import solve_ivp
import scipy as sci
import matplotlib as mpl
import tkinter as tk
import n_body_app
import os
import sys
plt.rc('mathtext', fontset="cm")

# %% Initialising all the planets, suns and objects that could be used in the simulation

"""Code to create and run GUI, from which you can change the configuration of the simulation"""

root = tk.Tk()
root.title("N Body Simulation")
root.geometry('400x300')
app = n_body_app.n_body_app(root)
root.mainloop()

planets = app.planets

print(planets)
objects = []

x = [Sun,  Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune, External_Planet
     , Butterfly_I, Butterfly_II, Butterfly_III, moth_I, moth_II, moth_III, bumblebee,
     pythag, pythag_I, Solar_System, Figure_8, goggles, yarn, yin_yang_I_a, Flower_in_circle, dragonfly
     ]

thisdict = {
    'Butterfly_I': Butterfly_I,
    'Butterfly_II': Butterfly_II,
    'Butterfly_III': Butterfly_III,
    'moth_I': moth_I,
    'moth_II': moth_II,
    'moth_III': moth_III,
    'bumblebee': bumblebee,
    'pythag': pythag,
    'pythag_I': pythag_I,
    'Solar_System': Solar_System,
    'Figure_8': Figure_8,
    'goggles': goggles,
    'yarn': yarn,
    'yin_yang_I_a': yin_yang_I_a,
    'Flower_in_circle': Flower_in_circle,
    'dragonfly': dragonfly
    }

try:
    for i in range(len(x)):
        if str(x[i].name) in planets:
            objects.append(x[i])
except AttributeError:
    if len(x) > 0:
        pass
    pass
    if len(planets) == 0:
        objects = thisdict['Solar_System']
        print("The defult configuration - the solar system - has been initiated")
    if len(planets) == 1:
        objects = thisdict[planets[0]]
    print(len(objects), " objects have been initialised into the simulation")

""" Defining the list of planets which will be used in the simulation, only the above objects can be placed in"""
""" Using the objects class to input all the initial variables and initiliase the planets """

solarsystem = SolarSystem(objects)

# %% Define constants
start=time.time()
"""
Defining the constant needed for the calculations. The initial positions and velocities needed to be converted
to m and m/s respectivly as the JPL NASA website provides it in AU and AU/Day. Other constants in clude the time
for the simulation which can be adjusted, and the number of iterations
"""

# Define universal gravitation constant
G = float(app.G.get())  # 6.67408e-11

# Reference quantities - to convert the emphemeris data from nasa to SI units
years = float(app.n_time_period.get())  # 100

# Define times
tStart = 0e0
time_period = float(app.time_period.get())  # 60*60*24*365
n_time_periods = years
iterations_year =  float(app.iterations.get()) # 25
iterations_total = n_time_periods*iterations_year
t_End = time_period*n_time_periods
max_steps = t_End/(iterations_total*4)
t=tStart
domain = (t, t_End)

# Non-Dimensionalization constants

K1=1
K2=1

initial = np.full((1, 6), 0, dtype=float)
mass = np.full((1, 1), 0, dtype=float)

name = ""

N = len(objects) # Find the number of objects in the Solar list

planets_initial = np.full([N, 6],0, dtype=float)

# Creates an array for all the particles used so that the initial positions and velocities are know
for i in range(len(solarsystem.planets)):
    planets_initial[i] = solarsystem.planets[i].init

planets_pos = planets_initial[:,0:3]
planets_vel = planets_initial[:,3:6]

# Create an array with all the masses
planets_mass = np.full((N, 1),0, dtype=float)

for i in range(len(solarsystem.planets)):
    planets_mass[i] = solarsystem.planets[i].mass

# Finds the centre of mass
moment_of_mass = np.array([0,0,0])
for i in range(N):
    moment_of_mass = moment_of_mass + planets_mass[i]*planets_pos[i]
r_com=(moment_of_mass)/(sum(planets_mass))
where_are_NaNs = np.isnan(r_com)
r_com[where_are_NaNs] = 0

# Finds centre of mass of velocity
momentum = np.array([0,0,0])
for i in range(N):
    momentum = momentum + planets_mass[i]*planets_vel[i]
v_com=(momentum)/(sum(planets_mass))
where_are_NaNs = np.isnan(v_com)
v_com[where_are_NaNs] = 0

# Flattens the array as thats the form solve_ivp takes it in

planets_pos = planets_pos.flatten()
planets_vel = planets_vel.flatten()
init_params=np.hstack((planets_pos, planets_vel))

# %% Solve the equation for Gravity for the n body system

# ## Get constantys from GUI

if len(app.ODE) > 1:
            print("Only one ODE solver must be selected!")
if len(app.ODE) == 0:
    app.ODE = ["RK45"]
    print("The defult ODE solver will be used as none were selected")

rtol = float(app.rtol.get())
atol = float(app.rtol.get())
integrator = app.ODE

# ## Run the solve_ivp solver
three_body_sol = sci.integrate.solve_ivp(fun=SolarSystem.ThreeBodyEquations, t_span=domain, y0=init_params, args=(
    G, planets_mass, N, K1, K2), max_step=max_steps, rtol=rtol, atol=atol, method=integrator[0])
t = three_body_sol['t']
iterations = len(three_body_sol['t']) # Find how many values of t were used

# ## Store the position solutions into three distinct arrays
r_sol = np.full((N*3,iterations),0)
r_sol = three_body_sol['y'][0:N*3,:]
r_sol = r_sol.T

# ## Chaanges the reference apoint to about that of the COM
momentum_com = np.full((iterations,3),0, dtype=float)

for i in range(N):
    momentum_com += planets_mass[i]*r_sol[:,i*3:(i+1)*3]

rcom_sol = momentum_com/sum(planets_mass)
where_are_NaNs = np.isnan(rcom_sol)
rcom_sol[where_are_NaNs] = 0
rearth_sol = r_sol[:,3:6]

r_com_sol = np.empty((iterations,N*3))

for i in range(N):
    r_com_sol[:,i*3:(i+1)*3] = r_sol[:,i*3:(i+1)*3] - rcom_sol

g = np.hstack((rcom_sol, rcom_sol,rcom_sol))

end=time.time()
print("Time for intialising data and integrating is: " , end-start)

# %% Analysis
start = time.time()
for planet in solarsystem.planets:
    planet.KE = np.full((1,1),0,dtype=float)
    planet.PE = np.full((1,1),0,dtype=float)
    planet.linear_m = np.full((1,3),0,dtype=float)
    planet.angular_m = np.full((1,1),0,dtype=float)

KE = []
PE = []
angular = []
linear = []
linear_x = []
linear_y = []
linear_z = []

colours = ['black','g','b','gold','y','m','c','r','lime','navy']

for col in range(iterations):
    sol = three_body_sol['y'][:,col]
    temp_KE, temp_PE, temp_angular, temp_linear, temp_linear_x, temp_linear_y, temp_linear_z = SolarSystem.getEnergy(
        solarsystem, sol, planets_mass, G, N, v_com)
    KE.append(temp_KE)
    PE.append(temp_PE)
    angular.append(temp_angular)
    linear.append(temp_linear)
    linear_x.append(temp_linear_x)
    linear_y.append(temp_linear_y)
    linear_z.append(temp_linear_z)

KE = np.array([KE])
PE = np.array([PE])
angular = np.array([angular])
linear = np.array([linear])
linear_x = np.array([linear_x])
linear_y = np.array([linear_y])
linear_z = np.array([linear_z])

"""Defining the energies which are the sum of other to plot"""
total = (KE+PE).flatten()

virial = (2*np.average(KE)+np.average(PE)).flatten()

save_results_to = sys.path[0] + "/Plots Generated/"

plotKE = plt.figure(1)
for planet in solarsystem.planets:
    i = solarsystem.planets.index(planet)
    plt.plot(t, planet.KE[1:,0], colours[i], label=solarsystem.planets[i].name, linewidth=0.9)
plt.title("Kinetic energy of individual planets", fontsize='9')
plt.xlabel(r"Time ($s$)")
plt.ylabel(r"Kinetic energy ($J$)")
plotKE.legend(loc='center right', bbox_to_anchor=(0.90, 0.5))
plotKE.show()
plotKE.savefig(save_results_to + 'planets_KE.png', dpi=600, bbox_inches='tight')

plotPE = plt.figure(2)
for planet in solarsystem.planets:
    i = solarsystem.planets.index(planet)
    plt.plot(t, planet.PE[1:,0], (colours)[i], label=solarsystem.planets[i].name, linewidth=0.9)
    plt.title("Potential energy of individual planets over $%.2f$ time periods" % (n_time_periods,),fontsize='9')
    plt.xlabel("Time ($s$)")
    plt.ylabel("Potential energy ($J$)")
    plotKE.legend()
    plotKE.show()
plotPE.show()
plotPE.savefig(save_results_to + "planets_PE.png", bbox_inches='tight')

plotTotal = plt.figure(3)
plt.plot(t, KE.flatten(), label="Kinetic energy",linewidth=0.9)
plt.plot(t, PE.flatten(), label="Potential energy", linewidth=0.9)
plt.plot(t, total, label="Total energy", linewidth=0.9)
plt.plot(t, np.full((1, len(t)), virial)[0,:], label="$KE + 2PE = 0$", linewidth=0.9)
plt.title("Comparison of KE, PE and total energy of objects over time",fontsize='9')
plt.xlabel("Time ($s$)")
plt.ylabel("Energy ($J$)")
plotTotal.legend(loc='center right', bbox_to_anchor=(0.90, 0.5))
plotTotal.show()
plotTotal.savefig(save_results_to +'Total_Energy_System.png', bbox_inches='tight', dpi=600)

plotOrbits = plt.figure(4)
for i in range(N):
    plt.plot(r_com_sol[:,i*3], r_com_sol[:,1+i*3], (colours)[i], label=solarsystem.planets[i].name,linewidth=0.9)
plt.title("Orbits mapped - 2D")
plt.xlabel("$x$ ($m$)")
plt.ylabel("$y$ ($m$)")
plotOrbits.legend(loc='center right', bbox_to_anchor=(0.90, 0.5))
plotOrbits.show()
plotOrbits.savefig(save_results_to +'Orbits_System.png', bbox_inches='tight', dpi=600)

plotLm = plt.figure(5)
plt.plot(t, linear.flatten(), linewidth=0.9)
plt.title("Total linear momentum of objects in the system over time", fontsize='9')
plt.xlabel("Time ($s$)")
plt.ylabel("Linear momentum ($kgms^{-2}$)")
plotLm.legend()
plotLm.show()
plotLm.savefig(save_results_to +'Linear_Momentum_System.png', bbox_inches='tight')

plotAm = plt.figure(6)
for planet in solarsystem.planets:
    i = solarsystem.planets.index(planet)
    plt.plot(t[1:], planet.angular_m[1:-1, 0], (colours)[i],
             label=solarsystem.planets[i].name, linewidth=0.9)
plt.plot(t, angular.flatten(), linewidth=0.9, label='Total')
plt.title("Total angular momentum of objects in the system over time", fontsize='9')
plt.xlabel("Time ($s$)")
plt.ylabel("Angular Momentum ($kgm^2s^{-1}$)")
plotAm.legend()
plotAm.show()
plotAm.savefig(save_results_to +'Angular_Momentum_System.png',dpi=600, bbox_inches='tight')

# threeD_plot = plt.figure(7)
threeD_plot = plt.figure(figsize=plt.figaspect(1)*2)
# ax = p3.Axes3D(threeD_plot)
ax = threeD_plot.gca(projection='3d', proj_type = 'ortho')
x_scale = 4
y_scale = 4
z_scale = 1
scale = np.diag([x_scale, y_scale, z_scale, 1.0])
scale = scale*(1.0/scale.max())
scale[3, 3] = 1.0

def short_proj():
  return np.dot(Axes3D.get_proj(ax), scale)

ax.get_proj = short_proj
ax.view_init(elev=45/2, azim=45)
ax.dist = 7
line = [ax.plot(r_com_sol[:, i*3], r_com_sol[:, i*3+1], r_com_sol[:,i*3+2], c=(colours)[i], label=solarsystem.planets[i].name, linewidth=0.9)[0] for i in range(N)]
ax.set_xlabel('')
ax.set_ylabel('')
ax.set_zlabel('')
ax.set_zlim([-0.5e12,0.5e12])
ax.set_title("Static 3D Orbit")
threeD_plot.legend(loc='center right')
threeD_plot.show()
threeD_plot.savefig(save_results_to +'3D static plot.png',dpi=600, bbox_inches='tight')
plt.clf()

plotLx = plt.figure(7)
for planet in solarsystem.planets:
    i = solarsystem.planets.index(planet)
    plt.plot(t, planet.linear_m[:-1, 0], (colours)[i], label=solarsystem.planets[i].name, linewidth=0.9)
plt.title("$x$ component of linear momentum over time", fontsize='9')
plt.xlabel("Time")
plt.ylabel("Linear momentum ($kgms^{-2}$)")
plotLx.legend(loc='upper left', bbox_to_anchor=(0.12, 0.80))
plotLx.show()
plotLx.show()
plotLx.savefig(save_results_to +'planets_momentum_x.png',dpi=600, bbox_inches='tight')

plotL_xyz = plt.figure(8)
plt.plot(t, linear_x.flatten(),linewidth=0.9, label='$x$')
plt.plot(t, linear_y.flatten(), linewidth=0.9, label='$y$')
plt.plot(t, linear_z.flatten(), linewidth=0.9, label='$z$')
plt.xlabel("Components of linear momentum over time", fontsize='9')
plt.ylabel("Linear momentum ($kgms^{-2}$)")
plotL_xyz.legend(loc='center left', bbox_to_anchor=(0.15, 0.75))
plotL_xyz.show()
plotL_xyz.savefig(save_results_to +'Linear_Momentum_System_components.png',dpi=600, bbox_inches='tight')

plot_planetsAngular = plt.figure(11)
for planet in solarsystem.planets:
    i = solarsystem.planets.index(planet)
    plt.plot(t[1:], planet.angular_m[1:-1, 0], (colours)[i],
             label=solarsystem.planets[i].name, linewidth=0.9)
plt.title("$x$ component of linear momentum over time", fontsize='9')
plt.xlabel("Time")
plt.ylabel("Angular momentum ($kgms^{-2}$)")
plot_planetsAngular.legend(loc='upper left', bbox_to_anchor=(0.12, 0.90))
plot_planetsAngular.show()
plot_planetsAngular.savefig(save_results_to + 'planets_momentum_angular.png', dpi=600, bbox_inches='tight')

end = time.time()

print("Time for calculating enrgy and plotting is:   " , end-start)

"""Crudely evaluates when the orbit and physics is no longer valid"""

total_off = ((total-total[0])/total)*100
total_cor = [i for i in total_off if i >= 1]

if len(total_cor) == 0:
    print("This orbit looks stable. The change in error from the beginiing to end is: ", total_off[-1],"%")
else:
    print("The stablility of the orbit seems to be only valid up till the ", len(total_cor), " time step")

# %% If animation wanted uncomment this section
"""
cm = plt.cm.get_cmap('tab10')
colours = cm.colors

anim_r_com_sol = r_com_sol[0::20,:].copy()
data_len = anim_r_com_sol.shape[:][0]

x = np.array(range(anim_r_com_sol.shape[:][0]))

# define new x range, we need 7 equally spaced values
xnew = np.linspace(x.min(), x.max(), 9)

# apply the interpolation to each column
f = interp1d(x, anim_r_com_sol, axis=0)

data = np.empty((data_len, N, 3))

for i in range(3):
    data[:,:,i] = anim_r_com_sol[:,i::3]

fig = plt.figure()
ax = p3.Axes3D(fig)
ax.set_xlabel('X')

ax.set_ylabel('Y')

ax.set_zlabel('Z')
ax.set_zlim([-3e12,3e12])

ax.set_title("Animated Orbit ")

h_particles = [ax.plot(*data[:1, i].T, marker='o', c=colours[i], ls='None')[0] for i in range(N)]
trace = [ax.plot(data[:1,i,0], data[:1,i,1], data[:1,i,2], c=colours[i])[0] for i in range(N)]

def update_particles(num):
    print(num, " out of ", len(x))
    global h_particles
    global trace

    for p in h_particles:
            p.remove()

    for t in trace:
        t.remove()

    trace = [ax.plot(data[:num,i,0], data[:num,i,1], data[:num,i,2], c=colours[i])[0] for i in range(N)]

    h_particles = [ax.plot(*data[num-1:num, i].T, marker='o', c=colours[i],label=solarsystem.planets[i].name, ls='None')[0] for i in range(N)]

    return h_particles, trace

prtcl_ani = animation.FuncAnimation(fig, update_particles, frames=data_len, interval=70, blit=False,repeat=False)

# writergif = animation.PillowWriter(fps=30)
# prtcl_ani.save('filename.gif', writer=writergif)
prtcl_ani.save(save_results_to + "Orbit_Animation.mp4", dpi=450)

# ['black','g','b','gold','y','m','c','r','lime']
"""

max_stars_repo_path: Simulation (GUI).py

max_stars_repo_name: hsspratt/Solar-system-nbody

max_issues_repo_path: Simulation (GUI).py

max_issues_repo_name: hsspratt/Solar-system-nbody

max_forks_repo_path: Simulation (GUI).py

max_forks_repo_name: hsspratt/Solar-system-nbody

reason: import numpy,import scipy,from scipy


text:
import re
import os
import json
import numpy as np
import xml.etree.ElementTree as ET

from typing import Dict

try:
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
except ImportError:
    pass

class WebotsGameLogParser:
    """
    The if the root class of the interface.
    It connects both the x3d and json parser and offers some high level example plots.
    """
    def __init__(self, log_folder: str, verbose=True):
        super().__init__()

        self.log_folder = log_folder

        x3d_file = [file for file in os.listdir(self.log_folder) if file.endswith(".x3d")][0]
        self.x3d = X3DParser(os.path.join(log_folder, x3d_file))

        game_json_file = [file for file in os.listdir(self.log_folder) if file.endswith(".json")][0]
        self.game_data = GameJsonParser(os.path.join(log_folder, game_json_file))

        if verbose:
            print(f"Duration: {self.get_max_player_timestamp() / 60 :.2f} Minutes")
            print(f"Players: {', '.join(self.x3d.get_player_names())}")

    def plot_ball_path(self):
        """
        Creates a combined plot with the paths for all the players
        """
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.set_xlim(-4.5, 4.5)
        ax.set_ylim(-3, 3)
        ax.axis('equal')
        ax.plot(*self.game_data.get_translations_for_id(self.x3d.get_ball_id()).T[0:2])
        plt.show()

    def plot_player_paths(self):
        """
        Creates a combined plot with the paths for all the players
        """
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.set_xlim(-4.5, 4.5)
        ax.set_ylim(-3, 3)
        ax.axis('equal')
        for bot in self.x3d.get_player_names():
            ax.plot(*self.game_data.get_translations_for_id(self.x3d.get_player_id(bot)).T[0:2])
        #ax.plot(*self.game_data.get_translations_for_id(self.x3d.get_ball_id()).T[0:2], linewidth=5)
        plt.show()

    def plot_path(self, id: int):
        """
        Creates a plot with the path for a certain object
        """
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.plot(*self.game_data.get_translations_for_id(id).T)
        plt.show()

    def get_max_player_timestamp(self) -> float:
        """
        Returns the last timestamp where any of the players has movement data
        """
        return max(self.game_data.get_timestamps_for_id(id).max() for id in self.x3d.get_player_ids())

class GameJsonParser:
    """
    Parses time series information from the .json log
    """
    def __init__(self, jsonfile: str):
        self.jsonfile = jsonfile

        # Load file contents
        with open(self.jsonfile, "r") as f:
            self.data = json.load(f)

        # Cache poses for id
        self._poses_for_id_cache = {}

    @property
    def get_time_step_size(self) -> int:
        """
        Returns the basic time step in ms.
        """
        return self.data["basicTimeStep"]

    def _parse_str_vector(self, vec: str) -> np.ndarray:
        """
        Converts a space divided string vector into a NumPy array.
        """
        return np.array([float(num) for num in vec.split(" ")], dtype=float)

    def get_poses_for_id(self, id: int) -> Dict[str, np.ndarray]:
        """
        Gets all the pose data for an object and
        return a list of dicts containing the time and the pose.
        """
        if not id in self._poses_for_id_cache.keys():
            poses = []
            for frame in self.data["frames"]:
                time = frame["time"]
                if "poses" in frame.keys():
                    for sim_object in frame["poses"]:
                        if sim_object["id"] == id:
                            if "translation" in sim_object.keys():
                                translation = self._parse_str_vector(sim_object["translation"])
                            if "rotation" in sim_object.keys():
                                rotation = self._parse_str_vector(sim_object["rotation"])
                            poses.append({
                                "time": time,
                                "trans": translation,
                                "rot": rotation
                            })
                            break
            # Clean out of bounds poses
            poses = self.cleanup_poses(poses)
            # Map the list of dicts to a dict of ndarrays to save memory
            self._poses_for_id_cache[id] = {key: np.array([p[key] for p in poses], dtype=float) for key in poses[0].keys()}

        return self._poses_for_id_cache[id]

    def get_translations_for_id(self, id: int) -> np.ndarray:
        """
        Returns an NumPy array with all translations for an object (in meters)
        """
        return self.get_poses_for_id(id)["trans"]

    def get_orientations_for_id(self, id: int) -> np.ndarray:
        """
        Returns an NumPy array with all orientations for an object (as axis angle)
        """
        return self.get_poses_for_id(id)["rot"]

    def get_timestamps_for_id(self, id: int) -> np.ndarray:
        """
        Returns an NumPy array with all timesteps for an object (in seconds)
        """
        # Get correct array and scale it to seconds
        return self.get_poses_for_id(id)["time"] / 1000

    def get_velocity_vectors_for_id(self, id: int):
        """
        Calcs vel vecs for a given id (in m/s)
        """
        Δtrans = np.diff(self.get_translations_for_id(id), axis=0)
        Δtime = np.diff(self.get_timestamps_for_id(id))
        return Δtrans / Δtime[:, np.newaxis]

    def get_velocities_for_id(self, id: int):
        """
        Calcs vels for a given id (in m/s)
        """
        return np.linalg.norm(self.get_velocity_vectors_for_id(id), axis=1)

    def cleanup_poses(self, poses: [dict])-> [dict]:
        """
        Removes out of bounds poses caused by e.g. the referee
        """
        def in_bounds(pose: dict, treshold: float = 4.5) -> bool:
            return all(treshold > np.abs(pose["trans"]))

        return list(filter(in_bounds, poses))

    def get_interpolated_translations(self, id: int, start: float = None, stop: float = None, step_size: float = 0.1) -> np.ndarray:
        """
        Returns the translation data interpolated to fit a given step size
        """
        # Get the data that should be interpolated
        fp = self.get_translations_for_id(id)
        xp = self.get_timestamps_for_id(id)
        # Add default values if necessary. Use the first and last timestep for the given id.
        if start is None: start = xp[0]
        if stop is None: stop = xp[-1]
        # Create the array of query values
        x_steps = np.arange(start=start, stop=stop, step=step_size)
        # Interpolate for each axis
        inter = np.vstack([np.interp(x_steps, xp, fp[:, i]) for i in range(3)])
        return inter.T

    def get_interpolated_orientations(self, id: int, start: float = None, stop: float = None, step_size: float = 0.1) -> np.ndarray:
        """
        Returns the orientation data interpolated to fit a given step size
        """
        # Get the data that should be interpolated
        fp = self.get_orientations_for_id(id)
        xp = self.get_timestamps_for_id(id)
        # Add default values if necessary. Use the first and last timestep for the given id.
        if start is None: start = xp[0]
        if stop is None: stop = xp[-1]
        # Create the array of query values
        x_steps = np.arange(start=start, stop=stop, step=step_size)
        # The orientation is in axes angles, so we need to split the angle in two components
        sin_angles = np.sin(fp[:, 3])
        cos_angles = np.cos(fp[:, 3])
        fp = [fp[:, 0], fp[:, 1], fp[:, 2], sin_angles, cos_angles]
        # Interpolate for each axis
        inter = np.vstack([np.interp(x_steps, xp, fpe) for fpe in fp])
        # Convert sin cos vector back to an single angle
        inter[3] = np.arctan2(inter[3], inter[4])
        return inter[:4].T

class X3DParser:
    """
    Parses the .x3d file in the gamelogs, that contains all the 3d object information.
    """
    def __init__(self, x3dfile: str):
        super(X3DParser).__init__()
        self.x3d_file_path = x3dfile

        # Load file contents
        with open(self.x3d_file_path, "r") as f:
            x3d_file_raw_content = f.read()

        self.xml_root = ET.fromstring(x3d_file_raw_content)

    def get_players(self) -> [dict]:
        """
        Get the names of the loaded players
        """
        def is_player(node: ET.Element) -> bool:
            return "player" in node.attrib.get("name", "")

        def simplify_dict(node: ET.Element) -> dict:
            return {"id": int(node.attrib["id"][1:]),
                    "name": str(node.attrib["name"])}

        return list(map(simplify_dict, filter(is_player, self.xml_root.iter("Transform"))))

    def get_player_names(self) -> [str]:
        """
        Returns a list with all player names
        """
        return list(map(lambda x: x["name"], self.get_players()))

    def get_player_id(self, name: str) -> int or None:
        """
        Returns the object id for a given player name
        """
        return (list(map(lambda x: x["id"], filter(lambda x: x["name"] == name, self.get_players()))) + [None])[0]

    def get_object_id(self, name: str) -> int or None:
        """
        Returns the object id for a object name
        """
        def equals_name(node: ET.Element) -> bool:
            return node.attrib.get("name", "") == name

        def simplify_to_id(node: ET.Element) -> id:
            return int(node.attrib["id"][1:])

        # Get all nodes from xml tree that match the name exactly
        nodes_filtered_by_name = filter(equals_name, self.xml_root.iter("Transform"))

        # Reduce them to their integer id
        ids_from_filtered_nodes = map(simplify_to_id, nodes_filtered_by_name)

        # Select the first candidate or None if there is no match
        first_or_none = (list(ids_from_filtered_nodes) + [None])[0]

        return first_or_none

    def get_player_ids(self) -> [int]:
        """
        Returns a list with all player object ids
        """
        return list(map(self.get_player_id, self.get_player_names()))

    def get_ball_id(self) -> int:
        """
        Returns the object id of the ball
        """
        return self.get_object_id("robocup soccer ball")

max_stars_repo_name: bit-bots/webots-web-log-interface

max_issues_repo_name: bit-bots/webots-web-log-interface

max_forks_repo_name: bit-bots/webots-web-log-interface


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

"""Analyse trained word context models."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import os

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from sklearn.manifold import TSNE

from word_context import configuration
from word_context import word_context_model
from word_context import tools

FLAGS = tf.flags.FLAGS

tf.flags.DEFINE_string("checkpoint_path", None,
                       "Path to checkpoint or the directory containing the "
                       "checkpoint to analyse.")

tf.flags.DEFINE_string("vocab_file", "",
                       "Text file mapping word to id.")

tf.flags.DEFINE_string("output_dir", "",
                       "Directory to store the output analysis to.")

tf.flags.DEFINE_string("sentences_file", "",
                       "Path to txt file containing sample sentences to "
                       "encode and visualize.")

tf.flags.DEFINE_boolean("plot", True,
                        "Whether to plot embeddings or not.")

tf.flags.DEFINE_string("mr_data_dir", "",
                       "Path to the Movie Review dataset. This should contain "
                       "the files rt-polarity.pos and rt-polarity.neg.")

# List of types to store output probabilities for.
TARGET_TYPES = ["the", "you", "alejo", "noon", "pablo", "gravel", "she",
                "decrepit", "conference"]

# List of tokens to compute the nearest neighbours to in embedding space.
SEED_TOKENS = ["<unk>", "the", "james", "least", "car", "red", "how",
               "organization"]

# List of test sentences for which we will compute norms.
TEST_SENTENCES = [
    "james",
    "car",
    "red car",
    "driving a red car on the road",
    "james was driving a red car on the road",
    "brain",
    "google brain",
    "bhuwan is doing an internship",
    "bhuwan is doing an internship at google brain",
    "sky",
    "blue sky",
    "beautiful",
    "such a beautiful",
    "such a beautiful blue sky today",
]

def _sq_euclidean_distance(x, y):
  """Compute squared euclidean distance between batches x and y.

  Args:
    x: 2D Numpy array.
    y: 2D Numpy array.

  Returns:
    distances: 2D Numpy array such that distances[i,j] =
      squared_L2_distance(x[i,:], y[j,:]).
  """
  sq_norm_x = (x ** 2).sum(axis=1)[:,None]
  sq_norm_y = (y ** 2).sum(axis=1)[None,:]
  inner_product = np.dot(x, np.transpose(y))
  distances = sq_norm_x + sq_norm_y - 2 * inner_product
  return np.clip(distances, 0., np.inf)

def _hyperbolic_distance(x, y):
  """Compute hyperbolic distance between batches x and y.

  Args:
    x: 2D Numpy array.
    y: 2D Numpy array.

  Returns:
    distances: 2D Numpy array such that distances[i,j] =
      hyperbolic(x[i,:], y[j,:]).
  """
  sq_norm_x = (x ** 2).sum(axis=1)[:,None]
  sq_norm_y = (y ** 2).sum(axis=1)[None,:]
  inner_product = np.dot(x, np.transpose(y))
  sq_norm_xy = sq_norm_x + sq_norm_y - 2 * inner_product
  return np.arccosh(1. + np.clip(
      2. * sq_norm_xy / np.clip((1. - sq_norm_x) * (1. - sq_norm_y), 1e-15, 1.),
      0., np.inf))

def _cosine_similarity(x, y):
  """Compute cosine similarity between batches x and y.

  Args:
    x: 2D Numpy array.
    y: 2D Numpy array.

  Returns:
    similarities: 2D Numpy array such that similarities[i, j] =
    cosine(x[i,:], y[j,:]).
  """
  normalized_x = x / np.linalg.norm(x, axis=1, keepdims=True)
  normalized_y = y / np.linalg.norm(y, axis=1, keepdims=True)

  return np.dot(normalized_x, np.transpose(normalized_y))

def _create_restore_fn(checkpoint_path, saver):
  """Creates a function that restores a model from checkpoint.

  Args:
    checkpoint_path: Checkpoint file or a directory containing a checkpoint
      file.
    saver: Saver for restoring variables from the checkpoint file.

  Returns:
    restore_fn: A function such that restore_fn(sess) loads model variables
      from the checkpoint file.
    checkpoint_path: Path to the model checkpoint that will be loaded

  Raises:
    ValueError: If checkpoint_path does not refer to a checkpoint file or a
      directory containing a checkpoint file.
  """
  if tf.gfile.IsDirectory(checkpoint_path):
    latest_checkpoint = tf.train.latest_checkpoint(checkpoint_path)
    if not latest_checkpoint:
      raise ValueError("No checkpoint file found in: %s" % checkpoint_path)
    checkpoint_path = latest_checkpoint

  def _restore_fn(sess):
    tf.logging.info("Loading model from checkpoint: %s", checkpoint_path)
    saver.restore(sess, checkpoint_path)
    tf.logging.info("Successfully loaded checkpoint: %s",
                    os.path.basename(checkpoint_path))

  return _restore_fn, checkpoint_path

def save_euclidean_norms(output_embeddings, vocab):
  """Compute euclidean norms and save in sorted order.

  Args:
    output_embeddings: Numpy array containing embeddings, one per row.
    vocab: List of tokens in the same order as they appear in output_embeddings.
  """
  norms = np.linalg.norm(output_embeddings, axis=1)

  sorted_idx = np.argsort(norms)
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "sorted_euclidean_norms.txt"),
                    "w")
  f.write("\n".join(["%s\t%.3f" % (vocab[ii], norms[ii])
                     for ii in sorted_idx]))
  f.close()

def save_hyperbolic_norms(output_embeddings, vocab):
  """Compute hyperbolic norms and save in sorted order.

  Args:
    output_embeddings: Numpy array containing embeddings, one per row.
    vocab: List of tokens in the same order as they appear in output_embeddings.
  """
  hyp_norms = np.squeeze(
      _hyperbolic_distance(output_embeddings,
                           np.zeros((1, output_embeddings.shape[1]))))

  sorted_idx = np.argsort(hyp_norms)
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "sorted_hyperbolic_norms.txt"),
                    "w")
  f.write("\n".join(["%s\t%.3f" % (vocab[ii], hyp_norms[ii])
                     for ii in sorted_idx]))
  f.close()

def save_nearest_neighbours(output_embeddings, vocab, seeds, num_nearest, metric):
  """Find nearest neighbours in embedding space to given seeds.

  Args:
    output_embeddings: Numpy array containing embeddings, one per row.
    vocab: List of tokens in the same order as they appear in output_embeddings.
    seeds: List of tokens to compute the nearest neighbours to.
    num_nearest: Number of neighbours to output.
    metric: "hyperbolic" or "cosine" or None.
  """
  seed_emb = np.zeros((len(seeds), output_embeddings.shape[1]))
  for ix, token in enumerate(vocab):
    if token in seeds:
      seed_emb[seeds.index(token), :] = output_embeddings[ix,:]

  if metric == "hyperbolic":
    distances = _hyperbolic_distance(seed_emb, output_embeddings)
  elif metric == "cosine":
    distances = - _cosine_similarity(seed_emb, output_embeddings)
  else:
    distances = _sq_euclidean_distance(seed_emb, output_embeddings)
  sorted_index = np.argsort(distances, axis=1)

  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "nearest_neighbours.txt"),
                    "w")
  for ix in range(sorted_index.shape[0]):
    f.write(seeds[ix] + " (%.3f)" % distances[ix, vocab.index(seeds[ix])] +
            " " + " ". join(["%s (%.3f)" % (vocab[ii], distances[ix,ii])
                             for ii in sorted_index[ix, :num_nearest]]) +
            "\n")
  f.close()

def build_hierarchy(thought_vectors, sentences, metric):
  """Recursively find the nearest neighbour with greater norm.

  Args:
    thought_vectors: 2D numpy array containing the thought vectors.
    sentences: Corresponding string sentences.
    metric: To compute distances.
  """
  if metric == "hyperbolic":
    distances = _hyperbolic_distance(thought_vectors, thought_vectors)
  elif metric == "cosine":
    distances = - _cosine_similarity(thought_vectors, thought_vectors)
  else:
    distances = _sq_euclidean_distance(thought_vectors, thought_vectors)

  sorted_index = np.argsort(distances, axis=1)

  hyp_norms = np.squeeze(
      _hyperbolic_distance(thought_vectors,
                           np.zeros((1, thought_vectors.shape[1]))))

  def _find_closest(n):
    """Return closest neighbour of n-th sentence with greater norm."""
    current_norm = hyp_norms[n]
    neighbour_norms = hyp_norms[sorted_index[n, 1:]]
    ii = 0
    while True:
      if ii == neighbour_norms.shape[0]: break
      if neighbour_norms[ii] > current_norm:
        break
      ii += 1
    if ii == neighbour_norms.shape[0]:
      return None
    else:
      return sorted_index[n, 1 + ii]

  f = tf.gfile.Open(os.path.join(FLAGS.output_dir,
                                 "sentence_hierarchies.txt"), "w")
  subset = np.random.randint(thought_vectors.shape[0], size=100)
  for i in range(100):
    current = subset[i]
    f.write(sentences[current] + " (%.4f)" % hyp_norms[current] + "\n")
    for j in range(10):
      current = _find_closest(current)
      if not current: break
      f.write(sentences[current] + " (%.4f)" % hyp_norms[current] + "\n")
    f.write("\n")
  f.close()

def elementwise_nearest_neighbours(thought_vectors, sentences, metric):
  """Find nearest neighbour to each given sentence.

  Args:
    thought_vectors: 2D numpy array containing the thought vectors.
    sentences: Corresponding string sentences.
    metric: To compute distances.
  """
  if metric == "hyperbolic":
    distances = _hyperbolic_distance(thought_vectors, thought_vectors)
  elif metric == "cosine":
    distances = - _cosine_similarity(thought_vectors, thought_vectors)
  else:
    distances = _sq_euclidean_distance(thought_vectors, thought_vectors)

  sorted_index = np.argsort(distances, axis=1)

  hyp_norms = np.squeeze(
      _hyperbolic_distance(thought_vectors,
                           np.zeros((1, thought_vectors.shape[1]))))

  f = tf.gfile.Open(os.path.join(FLAGS.output_dir,
                                 "sentence_nearest_neighbours.txt"), "w")
  for i in range(thought_vectors.shape[0]):
    f.write(sentences[i] + " (%.4f)" % hyp_norms[i] + "\n")
    for k in range(10):
      f.write(sentences[sorted_index[i,k+1]] +
              " (%.4f)" % hyp_norms[sorted_index[i,k+1]] + "\n")
    f.write("\n")
  f.close()

def save_test_sentence_norms(thought_vectors, sentences, metric, name):
  """Compute norms of the given thought vectors and save with the sentences.

  Args:
    thought_vectors: 2D numpy array containing the thought vectors in the same
      order as sentences.
    sentences: List of strings corresponding to the vectors above.
    metric: "hyperbolic" or "cosine" or None.
    name: File prefix for saving.
  """
  if metric == "hyperbolic":
    norms = np.squeeze(_hyperbolic_distance(
        thought_vectors, np.zeros((1, thought_vectors.shape[1]))))
  elif metric == "cosine":
    norms = np.linalg.norm(thought_vectors, axis=1)
  else:
    norms = np.linalg.norm(thought_vectors, axis=1)

  sorted_index = np.argsort(norms)

  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, name + "_sentence_norms.txt"),
                    "w")
  for ix in sorted_index:
    f.write(sentences[ix] + "\t" + "%.2f" % norms[ix] + "\n")
  f.close()

def save_target_probabilities(logits, targets, vocab):
  """Save the prediction probability of selected words.

  Args:
    logits: 2D Numpy array. Output logits per word.
    targets: 1D Numpy array. Target word to be predicted.
    vocab: Dict mapping words to their ids.
  """
  assert logits.shape[0] == targets.shape[0], (
      "Number of targets (%d) != Number of logits (%d)" % (targets.shape[0],
                                                           logits.shape[0]))
  types = []
  average_probabilities = []
  for typ in TARGET_TYPES:
    if typ not in vocab:
      tf.logging.info("%s OOV. Skipping.", typ)
      continue
    index = vocab[typ]

    occurrences = np.where(targets == index)[0]
    if occurrences.shape[0] == 0:
      tf.logging.info("No target matching %s. Skipping.", typ)
      continue

    prob_sum = 0.
    for ix in occurrences:
      scores = np.exp(logits[ix,:])
      probabilities = scores / np.sum(scores)
      prob_sum += probabilities[index]
    average_probabilities.append(prob_sum / occurrences.shape[0])
    types.append(typ)

  sorted_index = np.argsort(average_probabilities)

  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "target_probabilities.txt"),
                    "w")
  for ix in sorted_index:
    f.write(types[ix] + "\t" + "%.2f" % average_probabilities[ix] + "\n")
  f.close()

def plot_average_norm_by_length(embeddings, sentences, metric):
  """Plots the average embeddings norm for each length of sentences.

  Args:
    embeddings: 2D numpy array.
    sentences: List of string sentences corresponding to the embeddings.
    metric: "hyperbolic" or "cosine" or None.
  """
  if metric == "hyperbolic":
    norms = np.squeeze(_hyperbolic_distance(
        embeddings, np.zeros((1, embeddings.shape[1]))))
  elif metric == "cosine":
    norms = np.linalg.norm(embeddings, axis=1)
  else:
    norms = np.linalg.norm(embeddings, axis=1)
  lengths = [len(sentence.split()) for sentence in sentences]
  length_sum_norms = {}
  length_counts = {}
  for i in range(len(sentences)):
    if lengths[i] not in length_sum_norms:
      length_sum_norms[lengths[i]] = 0
      length_counts[lengths[i]] = 0
    length_sum_norms[lengths[i]] += norms[i]
    length_counts[lengths[i]] += 1

  tf.logging.info("Average Norm = %.4f",
                  float(sum(length_sum_norms.values())) /
                  sum(length_counts.values()))

  plot_x = []
  plot_y = []
  for length in length_sum_norms.keys():
    plot_x.append(length)
    plot_y.append(float(length_sum_norms[length]) / length_counts[length])

  fig, ax = plt.subplots()
  ax.plot(plot_x, plot_y)
  ax.set_xlabel('Sentence Length')
  ax.set_ylabel('Average Norm')
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "average_lengths.png"), "w")
  plt.savefig(f, dpi=1000)
  f.close()

def plot_embeddings(embeddings, labels, num_to_plot, name, metric, colors=None):
  """Plots a subset of the provided embeddings after projecting to 2D.

  If the embedding dimension > 2, these are first converted to 2 dimensions
  using T-SNE.

  Args:
    embeddings: 2D Numpy array containing embeddings one per row.
    labels: List of strings corresponding to annotations of the embeddings.
      If None no labels are annotated.
    num_to_plot: Number of embeddings to plot. These are randomly selected.
      Provide -1 to plot all embeddings
    name: Filename to save the plot in. .png extension will be appended.
    metric: Either "hyperbolic" or "cosine".

  Raises:
    ValueError: If the embeddings are less than 2 dimensions. A different
      visualization would be more appropriate in this case.
  """
  if embeddings.shape[1] < 2:
    raise ValueError("Embeddings must be at least 2 dimensional (found %d).",
                     embeddings.shape[1])
  elif embeddings.shape[1] == 2:
    tf.logging.info("Embeddings already 2 dimensional. No need for T-SNE.")
    embeddings_2d = embeddings
  else:
    tf.logging.info("Embeddings %d dimensional. Using T-SNE to project to 2 "
                    "dimensions.", embeddings.shape[1])
    model = TSNE(n_components=2, verbose=3, metric="precomputed")
    if metric == "hyperbolic":
      distances = _hyperbolic_distance(embeddings, embeddings)
    elif metric == "cosine":
      distances = - _cosine_similarity(embeddings, embeddings)
    else:
      distances = _sq_euclidean_distance(embeddings, embeddings)

    embeddings_2d = model.fit_transform(distances)
    tf.logging.info("Done.")

  if num_to_plot != -1:
    # Randomly sample subset to plot
    subset = np.random.randint(embeddings.shape[0], size=num_to_plot)
    embeddings_to_plot = embeddings_2d[subset, :]
    if colors:
      colors_to_plot = [colors[ii] for ii in subset]
    else:
      colors_to_plot = None
  else:
    subset = np.arange(embeddings.shape[0])
    embeddings_to_plot = embeddings_2d
    colors_to_plot = colors

  # Plot the subset
  fig, ax = plt.subplots()
  if colors_to_plot:
    ax.scatter(embeddings_to_plot[:,0], embeddings_to_plot[:,1],
               c=colors_to_plot)
  else:
    ax.scatter(embeddings_to_plot[:,0], embeddings_to_plot[:,1])
  if labels is not None:
    for i, ix in enumerate(subset):
      ax.annotate(labels[ix].decode(encoding="utf-8", errors="ignore"),
                  (embeddings_to_plot[i,0], embeddings_to_plot[i,1]),
                  fontsize=3)

  # Plot the unit circle for reference
  t = np.linspace(0, 2 * np.pi, 100)
  ax.plot(np.cos(t), np.sin(t), linewidth=1)

  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, name + ".png"),
                    "w")
  plt.savefig(f, dpi=1000)
  f.close()

def _create_decode_batch(sentences, vocab, length, pre):
  """Convert sentences into a batch for decoding.

  Args:
    sentences: List of strings to be decoded.
    vocab: A dict mapping tokens to indices.
    length: Minimum number of tokens to keep per instance.
    pre: Boolean indicating whether this batch corresponds to previous
      sentences or successor sentences.

  Returns:
    decode_ids: 2D padded numpy array containing token ids.
    decode_mask: 2D padded numpy array indicating where padding occurs.
  """
  token_lists = [sentence.lower().split() + ["<eos>"] for sentence in sentences]
  sentence_lengths = [len(sentence) for sentence in token_lists]

  decode_ids = np.zeros((len(token_lists), length), dtype="int32")
  decode_mask = np.zeros((len(token_lists), length), dtype="int32")
  for ix, sentence in enumerate(token_lists):
    to_fill = min(length, sentence_lengths[ix])
    for ii in range(to_fill):
      if pre:
        decode_ids[ix, -(ii+1)] = vocab.get(sentence[-(ii+1)], vocab["<unk>"])
        decode_mask[ix, -(ii+1)] = 1
      else:
        decode_ids[ix, ii] = vocab.get(sentence[ii], vocab["<unk>"])
        decode_mask[ix, ii] = 1

  return decode_ids, decode_mask

def _create_encode_batch(sentences, vocab):
  """Convert sentences into a batch.

  Args:
    vocab: A dict mapping tokens to indices.

  Returns:
    encode_ids: 2D padded numpy array containing token ids.
    encode_mask: 2D padded numpy array indicating where padding occurs.
  """
  token_lists = [sentence.lower().split() + ["<eos>"] for sentence in sentences]
  sentence_lengths = [len(sentence) for sentence in token_lists]

  encode_ids = np.zeros((len(token_lists), max(sentence_lengths)),
                        dtype="int32")
  encode_mask = np.zeros((len(token_lists), max(sentence_lengths)),
                         dtype="int32")
  for ix, sentence in enumerate(token_lists):
    encode_ids[ix, :sentence_lengths[ix]] = [vocab.get(token, vocab["<unk>"])
                                             for token in sentence]
    encode_mask[ix, :sentence_lengths[ix]] = 1

  return encode_ids, encode_mask

def _load_mr_dataset():
  pos_file = tf.gfile.Open(os.path.join(FLAGS.mr_data_dir, "rt-polarity.pos"))
  sentences = pos_file.read().splitlines()
  neg_file = tf.gfile.Open(os.path.join(FLAGS.mr_data_dir, "rt-polarity.neg"))
  sentences += neg_file.read().splitlines()
  return sentences

def main(unused_argv):
  # Create training directory if it doesn't already exist.
  if not tf.gfile.IsDirectory(FLAGS.output_dir):
    tf.logging.info("Creating training directory: %s", FLAGS.output_dir)
    tf.gfile.MakeDirs(FLAGS.output_dir)

  # Set up the model config.
  # Set up the model config.
  with tf.gfile.Open(
      os.path.join(FLAGS.checkpoint_path, "model_config.json")) as f:
    model_config_overrides = json.load(f)
  model_config = configuration.model_config(**model_config_overrides)
  tf.logging.info("model_config: %s",
                  json.dumps(model_config.values(), indent=2))

  # Read the vocabulary
  tf.logging.info("Reading vocabulary from: %s", FLAGS.vocab_file)
  vocab = tf.gfile.Open(FLAGS.vocab_file).read().splitlines()
  vocab_dict = {v: i for i,v in enumerate(vocab)}

  decode_length = model_config.context_window + model_config.condition_length

  tf.logging.info("Building training graph.")
  g = tf.Graph()
  with g.as_default():
    # Build the model.
    tf.logging.info("Building model.")
    model = word_context_model.WordContextModel(
        model_config, mode="decode")
    model.build()
    saver = tf.train.Saver()
    # Create restore function
    restore_model, ckpt = _create_restore_fn(FLAGS.checkpoint_path, saver)

  # Fetch input embeddings from the model
  skip_thought_emb = tools.load_skip_thought_embeddings(
      ckpt, model_config.vocab_size, model_config.word_embedding_dim)

  with tf.Session(graph=g) as sess:
    restore_model(sess)
    #sess.run(tf.global_variables_initializer())
    #for var in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
    #  print(var.name)
    #skip_thought_emb = sess.run("word_embedding:0")

    # Fetch output embeddings from the model
    output_embeddings = sess.run(model.output_embeddings)
    #thought_projection_matrix = sess.run("decoder/thoughts_projection/weights:0")
    #context_projection_matrix = sess.run("decoder/context_projection/weights:0")

    # Encode test sentences
    encode_ids, encode_mask = _create_encode_batch(TEST_SENTENCES, vocab_dict)
    encode_emb = skip_thought_emb[encode_ids, :]
    test_thought_vectors, test_thought_vectors_pre = sess.run(
        [model.thought_vectors, model.thought_vectors_pre],
        feed_dict = {model.encode_emb: encode_emb,
                     model.encode_mask: encode_mask})

    # Fetch local context vectors from batch
    batch_sentences = tf.gfile.Open(FLAGS.sentences_file).read().splitlines()
    encode_ids, encode_mask = _create_encode_batch(
        batch_sentences[1:-1], vocab_dict)
    decode_pre_ids, decode_pre_mask = _create_decode_batch(
        batch_sentences[:-2], vocab_dict, decode_length, True)
    decode_post_ids, decode_post_mask = _create_decode_batch(
        batch_sentences[2:], vocab_dict, decode_length, False)
    encode_emb = skip_thought_emb[encode_ids, :]
    decode_pre_emb = skip_thought_emb[decode_pre_ids, :]
    decode_post_emb = skip_thought_emb[decode_post_ids, :]
    if model_config.decode_strategy == "conditional":
      (batch_outputs, batch_thought_vectors, batch_thought_vectors_pre,
       context_vectors, batch_logits, batch_targets,
       batch_target_mask) = sess.run(
           [model.encoder_outputs, model.thought_vectors,
            model.thought_vectors_pre, model.local_context_vectors,
            model.logits, model.targets, model.target_mask],
           feed_dict = {model.encode_emb: encode_emb,
                        model.encode_mask: encode_mask,
                        model.decode_pre_ids: decode_pre_ids,
                        model.decode_pre_emb: decode_pre_emb,
                        model.decode_pre_mask: decode_pre_mask,
                        model.decode_post_ids: decode_post_ids,
                        model.decode_post_emb: decode_post_emb,
                        model.decode_post_mask: decode_post_mask})
    else:
      tf.logging.info("Skipping local context vector computation since "
                      "decode_strategy != conditional.")
      context_vectors = None

  # Analyse outputs
  tf.logging.info("Analysing output norms")
  save_hyperbolic_norms(output_embeddings, vocab)
  save_euclidean_norms(output_embeddings, vocab)

  tf.logging.info("Finding nearest neighbours")
  save_nearest_neighbours(output_embeddings, vocab, SEED_TOKENS, 10,
                          model_config.logit_metric)

  tf.logging.info("Computing test sentence norms")
  save_test_sentence_norms(test_thought_vectors, TEST_SENTENCES,
                           model_config.logit_metric, "test")
  save_test_sentence_norms(batch_thought_vectors, batch_sentences,
                           model_config.logit_metric, "batch")

  tf.logging.info("Finding sentence nearest neighbours")
  elementwise_nearest_neighbours(test_thought_vectors, TEST_SENTENCES,
                                 model_config.logit_metric)

  tf.logging.info("Building sentence hierarchies")
  build_hierarchy(batch_thought_vectors, batch_sentences,
                  model_config.logit_metric)

  tf.logging.info("Computing target word probabilities")
  batch_targets_reshaped = np.reshape(batch_targets, (encode_ids.shape[0], -1))
  batch_target_mask_reshaped = np.reshape(batch_target_mask, (encode_ids.shape[0], -1))
  save_target_probabilities(batch_logits, batch_targets, vocab_dict)

  tf.logging.info("Saving word embeddings")
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "word_embeddings.npy"), "w")
  np.save(f, output_embeddings)
  f.close()
  tf.logging.info("Saving output embeddings")
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "batch_outputs.npy"), "w")
  np.save(f, batch_outputs)
  f.close()
  tf.logging.info("Saving thought vectors")
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "batch_thoughts.npy"), "w")
  np.save(f, batch_thought_vectors)
  f.close()
  tf.logging.info("Saving thought vectors before scaling")
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "batch_thoughts_pre.npy"), "w")
  np.save(f, batch_thought_vectors_pre)
  f.close()
  tf.logging.info("Saving thought vectors projection matrix")
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "thoughts_projection.npy"), "w")
  np.save(f, thought_projection_matrix)
  f.close()
  tf.logging.info("Saving context vectors projection matrix")
  f = tf.gfile.Open(os.path.join(FLAGS.output_dir, "context_projection.npy"), "w")
  np.save(f, context_projection_matrix)
  f.close()

  if FLAGS.plot:
    tf.logging.info("Plotting length distribution")
    if model_config.logit_metric == "cosine":
      rescaled_vectors = batch_thought_vectors / np.linalg.norm(
          batch_thought_vectors, axis=1, keepdims=True)
    else:
      rescaled_vectors = batch_thought_vectors
    plot_average_norm_by_length(rescaled_vectors, batch_sentences,
                                model_config.logit_metric)

    tf.logging.info("Plotting word embeddings")
    if model_config.logit_metric == "hyperbolic":
      reparameterized_emb = output_embeddings
    elif model_config.logit_metric == "cosine":
      reparameterized_emb = output_embeddings / np.linalg.norm(
          output_embeddings, axis=1, keepdims=True)
    else:
      reparameterized_emb = output_embeddings
    plot_embeddings(reparameterized_emb, vocab, 1000, "word_embeddings",
                    model_config.logit_metric)

    tf.logging.info("Plotting batch thought vectors")
    if model_config.logit_metric == "hyperbolic":
      rescaled_vectors = batch_thought_vectors
    elif model_config.logit_metric == "cosine":
      rescaled_vectors = batch_thought_vectors / np.linalg.norm(
          batch_thought_vectors, axis=1, keepdims=True)
    else:
      rescaled_vectors = batch_thought_vectors
    plot_embeddings(rescaled_vectors, batch_sentences, -1,
                    "batch_thought_vectors",
                    model_config.logit_metric)

    tf.logging.info("Plotting test thought vectors")
    if model_config.logit_metric == "hyperbolic":
      rescaled_vectors = test_thought_vectors
    elif model_config.logit_metric == "cosine":
      rescaled_vectors = test_thought_vectors / np.linalg.norm(
          test_thought_vectors, axis=1, keepdims=True)
    else:
      rescaled_vectors = test_thought_vectors
    plot_embeddings(rescaled_vectors, TEST_SENTENCES, -1,
                    "test_thought_vectors",
                    model_config.logit_metric)

    tf.logging.info("Plotting test thought vectors and word embeddings")
    # Randomly sample subset to plot
    subset = np.random.randint(output_embeddings.shape[0], size=500)
    rescaled_vectors = np.vstack([test_thought_vectors,
                                  output_embeddings[subset,:]])
    labels = TEST_SENTENCES + [vocab[ix] for ix in subset]
    colors = [0] * len(TEST_SENTENCES) + [1] * subset.shape[0]
    if model_config.logit_metric == "cosine":
      rescaled_vectors = rescaled_vectors / np.linalg.norm(
          rescaled_vectors, axis=1, keepdims=True)
    plot_embeddings(rescaled_vectors, labels, -1,
                    "all_vectors",
                    model_config.logit_metric,
                    colors=colors)

    if model_config.decode_strategy == "conditional":
      tf.logging.info("Plotting context vectors")
      if model_config.logit_metric == "hyperbolic":
        rescaled_vectors = context_vectors
      elif model_config.logit_metric == "cosine":
        rescaled_vectors = context_vectors / np.linalg.norm(
            context_vectors, axis=1, keepdims=True)
      else:
        rescaled_vectors = context_vectors
      plot_embeddings(rescaled_vectors, None, -1, "context_vectors",
                      model_config.logit_metric)

if __name__ == '__main__':
  tf.app.run()

max_stars_repo_name: brain-research/hyperbolictext

max_issues_repo_name: brain-research/hyperbolictext

max_forks_repo_name: brain-research/hyperbolictext


text:
import athenaplotlib as apl

from athenaplotlib import Conv,Const

import argparse
from numba import njit,prange
import numpy as np
from scipy.stats import gaussian_kde

import mpl_scatter_density

import matplotlib.gridspec as gridspec

# Make the norm object to define the image stretch
from astropy.visualization import LogStretch
from astropy.visualization.mpl_normalize import ImageNormalize

import pandas as pd

import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib.cm import ScalarMappable

@njit
def linearPlotDataInWCR(r_array,quant_array,inWCR_array):
  """
  Build linear plot data for bottom graph,
  Outputs:
    WCR_array: Plottable array containing x data, r, and y data, the desired quantity
  """
  ncellsinwcr = np.count_nonzero(inWCR_array)
  plot_array  = np.zeros((2,ncellsinwcr))
  idx = 0
  for i in range(len(inWCR_array)):
    if inWCR_array[i] == True:
      plot_array[0,idx] = r_array[i]
      plot_array[1,idx] = quant_array[i]
      idx += 1
  return plot_array

@njit
def slicePlotDataInWCR(x_array,y_array,z_array,inWCR_array):
  """
  Ways to improve on this:
  include a translation array containing index number of element, allows for parallel transfer of memory and minimises reuse, but this is fine for now
  """
  ninWCR = np.count_nonzero(inWCR_array)
  slice_array = np.zeros((3,ninWCR))
  idx = 0
  for i in range(len(inWCR_array)):
    if inWCR_array[i] == True:
      slice_array[0,idx] = x_array[i]
      slice_array[1,idx] = y_array[i]
      slice_array[2,idx] = z_array[i]
      idx += 1
  return slice_array

def plotQuantity(grid_image_data,wcr_image_data,graph_data,config,quant):
  """
  Plot a quantity for this script, script

  inputs:
    array grid_image_data: data for upper plot, 2d slice with a quantity
    array wcr_image_data:  WCR data for upper plot, overlaid on top
    array graph_data:      distance/quantity plot data for lower graph
    dict  config:          config file dictionary
    str   quant:           quantity name, must be in all arrays
  """
  # Enable LaTeX plotting backend, requires TeXLive to be installed
  plt.rcParams.update({
      "text.usetex": True,
      "font.family": "serif",
  })

  fig,axs = plt.subplots(2)

  return

def processQuantity():
  #
  return

def processPlot(config,problem,data,base_filename):
  # Get simulation time
  time    = float(apl.readParameter(data,"Time",fallback=0.0))
  # Import star properties and calculate orbital position of stars
  star_WR = apl.Star(problem,"WR")
  star_OB = apl.Star(problem,"OB")
  star_WR.pos,star_OB.pos = apl.calcOrbit(star_WR,star_OB,time)

  # Determine where centerpoint co-ordinate for distance calculations is
  try:
    centerpoint = config["centerpoint"].lower()
    if   centerpoint == "apex":
      center_pos = apl.calcWCRApexPosition(star_WR,star_OB)
    elif centerpoint == "barycenter":
      center_pos = [0.0,0.0,0.0]
    elif centerpoint == "wr":
      center_pos = star_WR.pos
    elif centerpoint == "ob":
      center_pos = star_OB.pos
    else:
      print("Invalid centerpoint argument, defaulting to 0,0,0")
      center_pos = [0.0,0.0,0.0]
  except:
    print("Missing centerpoint argument in config file")
    raise

  # Reduce data from data file
  # Build list of quantities needed
  quantities = ["rho","vel1","vel2","vel3","r0"] # These values always needed
  for plot in config["plots"]:
    quantities.append(plot["quantity"])
  # Deduplicate
  quantities = list(dict.fromkeys(quantities))
  # Finally, reduce data
  reduced_data = apl.reduce(data,quantities,
                            avgmass_WR = star_WR.avgm,
                            avgmass_OB = star_OB.avgm)
  # Calculate distance from centerpoint for each cell
  wcr_method = apl.readParameter(config,"wcr_method",fallback="density")
  wcr_scale  = float(apl.readParameter(config,"wcr_scale",fallback=1.0))
  reduced_data["r"] = apl.calcDistFromPoint(reduced_data,center_pos)

  reduced_data["incwb"] = apl.detIfInWCR(reduced_data,
                                         star_WR, star_OB,
                                         scale=wcr_scale)

  xyscale = float(apl.readParameter(config,"xyscale",fallback=1.0))

  for plot in config["plots"]:
    # Build slice subset of data contained in
    wcr_linear = linearPlotDataInWCR(reduced_data["r"],
                                     reduced_data[plot["quantity"]],
                                     reduced_data["incwb"])
    wcr_slice  = slicePlotDataInWCR(reduced_data["xyz"][0],
                                    reduced_data["xyz"][1],
                                    reduced_data[plot["quantity"]],
                                    reduced_data["incwb"])
    # Build plotting environment
    cmap = apl.readParameter(config,"colourmap",fallback="viridis")
    # Use LaTeX environment
    plt.rcParams.update({
      "text.usetex": True,
      "font.family": "serif",
    })
    # Two figures being plotted
    fig = plt.figure(figsize=(5,8))
    gs  = gridspec.GridSpec(2, 1,height_ratios=[2,1])
    # gs = gridspec.GridSpec(1,1)
    slice_plot   = fig.add_subplot(gs[0])
    slice_plot.set_aspect("equal")
    apl.plotSlice(fig,slice_plot,
                  reduced_data["xyz"][0],
                  reduced_data["xyz"][1],
                  reduced_data[plot["quantity"]],
                  plot,
                  config)
    apl.showCWB(slice_plot,reduced_data,xyscale,alpha=0.5)

    vmin = apl.readParameter(plot,"zmin")
    vmax = apl.readParameter(plot,"vmax")
     # Should be able to move this to plotslice() with fig!

    scatter_plot = fig.add_subplot(gs[1],projection="scatter_density")
    apl.plotQuantVsDistScatter(scatter_plot,
                               wcr_linear[0],
                               wcr_linear[1],
                               plot,config)
    suptitle = apl.readParameter(config,"suptitle",fallback="")
    fig.suptitle(suptitle)
    # Finish up and save figure
    quant_name = apl.readParameter(plot, "filename", quit_on_fail=True)
    extension = apl.readParameter(config, "extension",fallback="png")
    dpi = int(apl.readParameter(config, "dpi",fallback=150))
    filename = "{0}-{1}.{2}".format(base_filename,quant_name,extension)
    plt.savefig(filename,dpi=300)

  # plot_array = linearPlotDataInWCR(data_reduced["d2"],
  #                                  data_reduced[quant],
                                  #  data_reduced["inWCR"])

  # r2arr = apl.calculateDistFromStar(data_red)
  # smooth_wind_density = apl.calculateSingleWindDensity(r2arr,mdot1,vinf1)
  # true_index = apl.detInWCR(data_red[3],smooth_wind_density)

  # plotarr = apl.buildPlotArray(data_red[3],r2arr,true_index)

  # norm = ImageNormalize(vmin=1., vmax=1000., stretch=LogStretch())
  # fig = plt.figure()
  # ax = fig.add_subplot(1,1,1,projection="scatter_density")
  # density = ax.scatter_density(plot_array[0],plot_array[1],norm=norm)
  # fig.colorbar(density)
  # ax.set_yscale("log")
  # plt.show()

    # Build plot array for the

  # plt.hist2d(plotarr[0],plotarr[1],(100,100),cmap=plt.cm.jet,norm=LogNorm())
  # plt.scatter(plotarr[0],plotarr[1],alpha=0.1)
  # plt.yscale("log")
  # plt.show()

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Process 2D and hst data from a folder of Athena outputs.\nRequires a config file and a positional argument of at least one config file, additional prefixes and directories can also be added")
  parser.add_argument("config_filename",
                      help="Name of config file used for this program")
  parser.add_argument("athena_input_filename",
                      help="Name of Athena config file, needed for config")
  parser.add_argument("HDF5_filename",
                      help="Name of HDF5 grid file, extension .athdf")
  parser.add_argument("output_filename",
                      help="Filename of output, will be stored in the form <output_filename>-<quantity>.<extension>")
  args,overrides = parser.parse_known_args()
  config  = apl.readConfigFile(args.config_filename)
  problem = apl.readProblemFile(args.athena_input_filename)
  data    = apl.readDataFile(args.HDF5_filename)
  output_filename = args.output_filename
  # Override configs
  if len(overrides) > 0:
    print("! Overriding config file with program arguments")
    for arg in overrides:
      try:
        key,val = arg.split("=",1)
        config[key] = val
        print("> {0} = \"{1}\"".format(key,val))
      except:
        pass
    print("! Done!")
  processPlot(config,problem,data,output_filename)

max_stars_repo_name: atomsite/ath-dust-plot

max_issues_repo_name: atomsite/ath-dust-plot

max_forks_repo_name: atomsite/ath-dust-plot

reason: import numpy,from scipy,from numba,from astropy


text:
###------------------------------------------------------------------------
###
### Program to copute interatomic distances from a non-orthonomal basis
###
###------------------------------------------------------------------------

# Import relevant libraries

import numpy as np
import os

# Lattice parameters - unit cell length and angles
a = 9.503044
b = 5.52558
c = 5.59397

al = np.deg2rad(89.826400)
be = np.deg2rad(125.757286)
ga = np.deg2rad(89.927844)

# Atomic positions
atomA = np.array([[0, 0, 0]]) # [0, 0, 0] for vector origin
atomB = np.array([[0, 0, 1]]) # position of iterest (fractional)
atomC = np.array([[0, 1, 0]]) # 2nd position of interest (mainly for angle calcs)

# Vector or  positions relative to atomA

vBA = atomB - atomA
vCA = atomC - atomA
vBC = atomB - atomC

# Triclinic Metric Tensor
g=np.array([[a ** 2, a * b * np.cos(ga), a * c * np.cos(be)], \
[b * a * np.cos(ga),b ** 2 , b * c * np.cos(al)], \
[c * a * np.cos(be), c * b * np.cos(al), c ** 2]])

# Orthonormal distance between atoms (matrix multiplication)
dBA = (vBA @ g @ vBA.transpose())**0.5
dCA = (vCA @ g @ vCA.transpose())**0.5
dBC = (vBC @ g @ vBC.transpose())**0.5

# dot product between B and C (matrix multiplication)
dpBC = vBA @ g @ vCA.transpose()

print('Metric tensor', g, sep = '\n')
print('Distance B to A: ',str(dBA).replace('[','').replace(']',''))
print('Distance C to A: ',str(dCA).replace('[','').replace(']',''))
print('Distance B to C: ',str(dBC).replace('[','').replace(']',''))

if np.array_equal(atomA, atomB) or np.array_equal(atomA, atomC):
    angleCAB = 'B or C cannot equal A!'
else:
    angleCAB = np.rad2deg(np.arccos(dpBC / (((dBA ** 2) * (dCA ** 2)) ** 0.5)))

print('Angle between B and C: ',str(angleCAB).replace('[','').replace(']',''))

max_stars_repo_name: antgobar/rmcalyse

max_issues_repo_name: antgobar/rmcalyse

max_forks_repo_name: antgobar/rmcalyse


text:
import numpy as np
import tensorflow as tf

from ..data import fields
from . import evaluator
from . import visualization

class SemSegEvaluator(evaluator.Evaluator):
    """Class to evaluate semantic segmentation metrics."""

    def __init__(
        self,
        category_index,
        ignore_label=255,
        include_metrics_per_category=False,
        max_examples_to_draw=10,
    ):
        """
        Args:
            category_index: a dict containing COCO-like category information keyed
                by the 'id' field of each category.
            ignore_label (int): value in semantic segmentation ground truth.
                Predictions for the corresponding pixels should be ignored.
        """
        super(SemSegEvaluator, self).__init__(category_index)
        self._ignore_label = ignore_label
        self._include_metrics_per_category = include_metrics_per_category
        self._num_classes = max(category_index.keys()) + 1
        self._N = self._num_classes + 1
        self._conf_matrix = np.zeros((self._N, self._N), dtype=np.int64)
        self._groundtruths = {}
        self._image_ids = {}
        self._summaries = []
        self._max_examples_to_draw = max_examples_to_draw

    def clear(self):
        self._conf_matrix = np.zeros((self._N, self._N), dtype=np.int64)
        self._groundtruths.clear()
        self._image_ids.clear()
        self._summaries = []

    def add_single_ground_truth_image_info(self, image_id, groundtruth_dict):
        """Adds groundtruth for a single image to be used for evaluation.
        Args:
            image_id: A unique string/integer identifier for the image.
            groundtruth_dict: A dictionary containing -
                fields.InputFields.sem_seg: unint8 numpy array of shape [height, width]
                    with values ranging between 0 ~ num_classes.
        Raises:
            ValueError: On adding groundtruth for an image more than once.
                Will also raise error if evaluate_masks and instance masks are not in
                groundtruth dictionary.
        """
        if image_id in self._image_ids:
            tf.logging.warning(
                'image %s has already been added to the ground truth database.',
                image_id)
            return
        sem_seg = groundtruth_dict[fields.InputFields.sem_seg]
        sem_seg[sem_seg == self._ignore_label] = self._num_classes

        if len(self._image_ids) < self._max_examples_to_draw:
            image = visualization.visualize_sem_seg_on_image_array(
                image=groundtruth_dict[fields.InputFields.image],
                sem_seg=sem_seg,
                category_index=self._category_index
            )
            self._summaries.append(
                tf.Summary.Value(
                    tag="{}/Groudtruth/Semantic".format(image_id),
                    image=tf.Summary.Image(
                        encoded_image_string=visualization.encode_image_array_as_png_str(image)))
            )
        self._groundtruths[image_id] = sem_seg
        self._image_ids[image_id] = False

    def add_single_predicted_image_info(self, image_id, result_dict):
        """Adds detections for a single image to be used for evaluation.
        Args:
        image_id: A unique string/integer identifier for the image.
        result_dict: A dictionary containing -
            fields.ResultFields.sem_seg: uint8 numpy array of shape [height, width]
                with values ranging between 0 ~ num_classes.
        Raises:
        ValueError: If detection masks are not in detections dictionary.
        """
        if image_id not in self._image_ids:
            raise ValueError('Missing groundtruth for image id: {}'.format(image_id))

        if self._image_ids[image_id]:
            tf.logging.warning(
                'image %s has already been added to the detection result database',
                image_id)
            return

        pred = result_dict[fields.ResultFields.sem_seg]

        if len(self._image_ids) <= self._max_examples_to_draw:
            image = visualization.visualize_sem_seg_on_image_array(
                image=result_dict[fields.InputFields.image],
                sem_seg=pred,
                category_index=self._category_index,
            )
            self._summaries.append(
                tf.Summary.Value(
                    tag="{}/Prediction/Semantic".format(image_id),
                    image=tf.Summary.Image(
                        encoded_image_string=visualization.encode_image_array_as_png_str(image)))
            )

        gt = self._groundtruths.pop(image_id)
        self._conf_matrix += np.bincount(
            self._N * pred.reshape(-1) + gt.reshape(-1), minlength=self._N ** 2
        ).reshape(self._N, self._N)
        self._image_ids[image_id] = True

    def evaluate(self):
        """Compute evaluation result.
        Returns:
        A dictionary of metrics with the following fields -
            'SemSegPerformance/mIOU',
            'SemSegPerformance/fwIOU',
            'SemSegPerformance/mACC',
            'SemSegPerformance/pACC'
            if include_metrics_per_category, category specific results with keys
            of the form -
                'SemSegPerformanceByCategory/IOU/category',
                'SemSegPerformanceByCategory/ACC/category'.
        """
        acc = np.zeros(self._num_classes, dtype=np.float)
        iou = np.zeros(self._num_classes, dtype=np.float)
        tp = self._conf_matrix.diagonal()[:-1].astype(np.float)
        pos_gt = np.sum(self._conf_matrix[:-1, :-1], axis=0).astype(np.float)
        class_weights = pos_gt / np.sum(pos_gt)
        pos_pred = np.sum(self._conf_matrix[:-1, :-1], axis=1).astype(np.float)
        acc_valid = pos_gt > 0
        acc[acc_valid] = tp[acc_valid] / pos_gt[acc_valid]
        iou_valid = (pos_gt + pos_pred) > 0
        union = pos_gt + pos_pred - tp
        iou[acc_valid] = tp[acc_valid] / union[acc_valid]

        metrics = {}
        metrics['SemSegPerformance/mIOU'] = np.sum(acc) / np.sum(acc_valid)
        metrics['SemSegPerformance/fwIOU'] = np.sum(iou) / np.sum(iou_valid)
        metrics['SemSegPerformance/mACC'] = np.sum(iou * class_weights)
        metrics['SemSegPerformance/pACC'] = np.sum(tp) / np.sum(pos_gt)

        if self._include_metrics_per_category:
            for category_id in self._category_index:
                category_name = self._category_index[category_id]["name"]
                metrics['SemSegPerformanceByCategory/IOU/{:s}'.format(
                    category_name)] = iou[category_id]
                metrics['SemSegPerformanceByCategory/ACC/{:s}'.format(
                    category_name)] = acc[category_id]

        for key, value in metrics.items():
            self._summaries.append(
                tf.Summary.Value(tag=key, simple_value=value)
            )
        return metrics, self._summaries


text:
# -*- coding: utf-8 -*-

"""

Author
------
Bo Zhang

Email
-----
bozhang@nao.cas.cn

Created on
----------
- Sat Jul 16 22:30:00 2016

Modifications
-------------
- Wed Jul 20 10:00:00 2016    creation

Aims
----
- scripts to run STARLIGHT_WRAPPER

"""

from __future__ import print_function

import os
import numpy as np
from astropy.table import Table, Column
from astropy.io import fits
# from bopy.spec.read_spectrum import read_spectrum

from .config import PACKAGE_PATH

# TODO: implement some template scripts

def ezrun():
    pass

def run():
    pass

reason: import numpy,from astropy


text:
from cppdefs import *
import numpy as np

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#
# # ROUTINE: ANALYTICAL_PROFILE
#
# DESCRIPTION:  This routine creates a vertical profile {\tt prof} with value
#               {\tt v1} in a surface layer down to depth {\tt z1} and a bottom
#               layer of value {\tt v2} reaching from depth {\tt z2} down to the bottom.
#               Both layers are connected by an intermediate layer reaching from {\tt z1}
#               to {\tt z2} with values linearly varying from {\tt v1} to {\tt v2}.
#
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

def analytical_profile(nlev, vertical_layers, depth1, surface_layer_value, depth2, bottom_layer_value):

    """
    Description: Creates a vertical profile with a value (surface_layer_value) in a
                 surface layer down to depth (depth1) and a bottom layer with a value
                 (bottom_layer_value) from depth (depth2) down to the bottom.
                 Both layers are connected by an intermediate layer reaching from
                 depth1 to depth2 with values varying linearly from surface_layer_value
                 to bottom_layer_value.

    :param nlev:
    :param vertical_layers:
    :param depth1:
    :param surface_layer_value: value at the surface
    :param depth2:
    :param bottom_layer_value: value at bottom
    :return: vertical profile
    """

    # OUTPUT PARAMETERS
    vertical_profile = np.empty(nlev,dtype=float)

    # LOCAL VARIABLES
    alpha = float()

    # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

    if (depth2 - depth1) > -1E-15:
        alpha = (bottom_layer_value - surface_layer_value) / (depth2 - depth1 + 2E-15)

    else:
        STDERR('*************************************************')
        STDERR('*  Error detected by analytical_profile.py:     *')
        STDERR('*  anz2 should be larger than anz1.             *')
        STDERR('*  Please edit BFM_General.nml or bio_bfm.nml.  *')
        STDERR('*************************************************')

    for i in range(nlev-1,0,-1):
        if vertical_layers[i] < depth1:
            vertical_profile[i] = surface_layer_value

        if alpha < 1E15:
            if depth1 < vertical_layers[i] < depth2:
                vertical_profile[i] = surface_layer_value + alpha * (vertical_layers[i] - depth1)

        if vertical_layers[i] > depth2:
            vertical_profile[i] = bottom_layer_value

    return vertical_profile

# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
#   Copyright 2013 BFM System Team (bfm_st@lists.cmcc.it)
#   Copyright by the GOTM-team under the GNU Public License - www.gnu.org
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=


text:
import logging
from typing import List, Tuple, Union

import numpy as np
from pymatgen import Structure
from pymatgen.io.ase import AseAtomsAdaptor
from spglib import spglib

from amset.constants import defaults, ktol

__author__ = "Alex Ganose"
__maintainer__ = "Alex Ganose"
__email__ = "aganose@lbl.gov"

logger = logging.getLogger(__name__)

def kpoints_to_first_bz(
    kpoints: np.ndarray, tol=ktol, negative_zone_boundary: bool = True
) -> np.ndarray:
    """Translate fractional k-points to the first Brillouin zone.

    I.e. all k-points will have fractional coordinates:
        -0.5 <= fractional coordinates < 0.5

    Args:
        kpoints: The k-points in fractional coordinates.
        tol: Fractional tolerance for evaluating zone boundary points.
        negative_zone_boundary: Whether to use -0.5 (spglib convention) or
            0.5 (VASP convention) for zone boundary points.

    Returns:
        The translated k-points.
    """
    kp = kpoints - np.round(kpoints)

    # account for small rounding errors for 0.5
    round_dp = int(np.log10(1 / tol))
    krounded = np.round(kp, round_dp)

    if negative_zone_boundary:
        kp[krounded == 0.5] = -0.5
    else:
        kp[krounded == -0.5] = 0.5
    return kp

def sort_boltztrap_to_spglib(kpoints):
    sort_idx = np.lexsort(
        (
            kpoints[:, 2],
            kpoints[:, 2] < 0,
            kpoints[:, 1],
            kpoints[:, 1] < 0,
            kpoints[:, 0],
            kpoints[:, 0] < 0,
        )
    )
    boltztrap_kpoints = kpoints[sort_idx]

    sort_idx = np.lexsort(
        (
            boltztrap_kpoints[:, 0],
            boltztrap_kpoints[:, 0] < 0,
            boltztrap_kpoints[:, 1],
            boltztrap_kpoints[:, 1] < 0,
            boltztrap_kpoints[:, 2],
            boltztrap_kpoints[:, 2] < 0,
        )
    )
    return sort_idx

def get_kpoints_tetrahedral(
    kpoint_mesh: Union[float, List[int]],
    structure: Structure,
    symprec: float = defaults["symprec"],
    time_reversal_symmetry: bool = True,
) -> Tuple[np.ndarray, ...]:
    """Gets the symmetry inequivalent k-points from a k-point mesh.

    Follows the same process as SpacegroupAnalyzer.get_ir_reciprocal_mesh
    but is faster and allows returning of the full k-point mesh and mapping.

    Args:
        kpoint_mesh: The k-point mesh as a 1x3 array. E.g.,``[6, 6, 6]``.
            Alternatively, if a single value is provided this will be
            treated as a k-point spacing cut-off and the k-points will be generated
            automatically.  Cutoff is length in Angstroms and corresponds to
            non-overlapping radius in a hypothetical supercell (Moreno-Soler length
            cutoff).
        structure: A structure.
        symprec: Symmetry tolerance used when determining the symmetry
            inequivalent k-points on which to interpolate.
        time_reversal_symmetry: Whether the system has time reversal symmetry.

    Returns:
        The irreducible k-points and their weights as tuple, formatted as::

            (ir_kpoints, weights)

        If return_full_kpoints, the data will be returned as::

            (ir_kpoints, weights, kpoints, ir_kpoints_idx, ir_to_full_idx)

        Where ``ir_kpoints_idx`` is the index of the unique irreducible k-points
        in ``kpoints``. ``ir_to_full_idx`` is a list of indices that can be
        used to construct the full Brillouin zone from the ir_mesh. Note the
        ir -> full conversion will only work with calculated scalar properties
        such as energy (not vector properties such as velocity).
    """
    from amset.electronic_structure.tetrahedron import get_tetrahedra

    if isinstance(kpoint_mesh, (int, float)):
        kpoint_mesh = get_kpoint_mesh(structure, kpoint_mesh)

    atoms = AseAtomsAdaptor().get_atoms(structure)

    if not symprec:
        symprec = 1e-8

    grid_mapping, grid_address = spglib.get_ir_reciprocal_mesh(
        kpoint_mesh, atoms, symprec=symprec, is_time_reversal=time_reversal_symmetry
    )
    full_kpoints = grid_address / kpoint_mesh

    tetra, ir_tetrahedra_idx, ir_tetrahedra_to_full_idx, tet_weights = get_tetrahedra(
        structure.lattice.reciprocal_lattice.matrix,
        grid_address,
        kpoint_mesh,
        grid_mapping,
    )

    ir_kpoints_idx, ir_to_full_idx, weights = np.unique(
        grid_mapping, return_inverse=True, return_counts=True
    )
    ir_kpoints = full_kpoints[ir_kpoints_idx]

    return (
        ir_kpoints,
        weights,
        full_kpoints,
        ir_kpoints_idx,
        ir_to_full_idx,
        tetra,
        ir_tetrahedra_idx,
        ir_tetrahedra_to_full_idx,
        tet_weights,
    )

def get_kpoint_mesh(structure: Structure, cutoff_length: float, force_odd: bool = True):
    """Calculate reciprocal-space sampling with real-space cut-off."""
    reciprocal_lattice = structure.lattice.reciprocal_lattice_crystallographic

    # Get reciprocal cell vector magnitudes
    abc_recip = np.array(reciprocal_lattice.abc)

    mesh = np.ceil(abc_recip * 2 * cutoff_length).astype(int)

    if force_odd:
        mesh += (mesh + 1) % 2

    return mesh

def get_mesh_from_kpoint_diff(kpoints):
    kpoints = np.array(kpoints)

    # whether the k-point mesh is shifted or Gamma centered mesh
    is_shifted = np.min(np.linalg.norm(kpoints, axis=1)) > 1e-6

    unique_a = np.unique(kpoints[:, 0])
    unique_b = np.unique(kpoints[:, 1])
    unique_c = np.unique(kpoints[:, 2])

    if len(unique_a) == 1:
        na = 1
    else:
        na = 1 / np.min(np.diff(unique_a))

    if len(unique_b) == 1:
        nb = 1
    else:
        nb = 1 / np.min(np.diff(unique_b))

    if len(unique_c) == 1:
        nc = 1
    else:
        nc = 1 / np.min(np.diff(unique_c))

    # due to limited precission of the input k-points, the mesh is returned as a float
    return np.array([na, nb, nc]), is_shifted

def get_mesh_from_kpoint_numbers(kpoints, tol=ktol):
    round_dp = int(np.log10(1 / tol))

    kpoints = kpoints_to_first_bz(kpoints)
    round_kpoints = np.round(kpoints, round_dp)

    nx = len(np.unique(round_kpoints[:, 0]))
    ny = len(np.unique(round_kpoints[:, 1]))
    nz = len(np.unique(round_kpoints[:, 2]))

    return nx, ny, nz

def get_kpoint_indices(kpoints, mesh, is_shifted=False):
    mesh = np.array(mesh)
    shift = np.array([1, 1, 1]) if is_shifted else np.array([0, 0, 0])
    min_kpoint = -np.floor(mesh / 2).round().astype(int)
    addresses = ((kpoints + shift / (mesh * 2)) * mesh).round().astype(int)
    shifted = addresses - min_kpoint
    nyz = mesh[1] * mesh[2]
    nz = mesh[2]
    indices = shifted[:, 0] * nyz + shifted[:, 1] * nz + shifted[:, 2]
    return indices.round().astype(int)

def get_kpoints_from_bandstructure(bandstructure, cartesian=False):
    if cartesian:
        return np.array([k.coords for k in bandstructure.kpoints])
    else:
        return np.array([k.frac_coords for k in bandstructure.kpoints])

max_stars_repo_name: kbspooner/amset

max_issues_repo_name: kbspooner/amset

max_forks_repo_name: kbspooner/amset


text:
from unittest.mock import ANY, call, Mock

import pytest
import numpy as np
import torch
import torch.nn as nn

from skorch.net import NeuralNet
from mltome.skorch.callbacks import LRRecorder
from mltome.skorch.callbacks import TensorboardXLogger

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(1, 2)
        self.layer2 = nn.Linear(2, 1)

    def forward(self, x):
        return self.layer2(self.layer1(x))

@pytest.fixture
def data():
    X = np.ones((1000, 1), dtype='float32')
    y = np.ones((1000, 1), dtype='float32')
    return X, y

# LRRecorder
def test_lr_recorder_default_lr_per_epoch(data):
    module = MyModule()
    X, y = data

    recorder = LRRecorder(per_epoch=True)
    net = NeuralNet(
        module,
        torch.nn.MSELoss,
        callbacks=[recorder],
        max_epochs=5,
        optimizer__lr=0.1,
        batch_size=100)
    net.fit(X, y)

    lrs = net.history[:, 'default_lr']
    assert lrs == [0.1] * 5

def test_lr_recorder_default_lr_per_batch(data):
    module = MyModule()
    X, y = data

    recorder = LRRecorder(per_epoch=False)
    net = NeuralNet(
        module,
        torch.nn.MSELoss,
        callbacks=[recorder],
        max_epochs=5,
        optimizer__lr=0.1,
        batch_size=100)
    net.fit(X, y)

    lrs_per_epoch = net.history[:, 'batches', :, 'default_lr']
    for lr in lrs_per_epoch:
        assert lr == [0.1] * 10

def test_lr_recorder_custom_lr_per_epoch(data):
    module = MyModule()
    X, y = data

    param_group = [('layer1.*', {'lr': 0.01})]
    params_names = [name for name, _ in param_group]
    recorder = LRRecorder(group_names=params_names, per_epoch=True)

    net = NeuralNet(
        module,
        torch.nn.MSELoss,
        callbacks=[recorder],
        max_epochs=5,
        optimizer__lr=0.1,
        batch_size=100,
        optimizer__param_groups=param_group)
    net.fit(X, y)

    lrs = net.history[:, 'layer1.*']
    assert lrs == [0.01] * 5

    lrs = net.history[:, 'default_lr']
    assert lrs == [0.10] * 5

def test_lr_recorder_custom_lr_per_batch(data):
    module = MyModule()
    X, y = data

    param_group = [('layer1.*', {'lr': 0.01})]
    params_names = [name for name, _ in param_group]
    recorder = LRRecorder(group_names=params_names, per_epoch=False)

    net = NeuralNet(
        module,
        torch.nn.MSELoss,
        callbacks=[recorder],
        max_epochs=5,
        optimizer__lr=0.1,
        batch_size=100,
        optimizer__param_groups=param_group)
    net.fit(X, y)

    lrs_per_epoch_layer1 = net.history[:, 'batches', :, 'layer1.*']
    for lr in lrs_per_epoch_layer1:
        assert lr == [0.01] * 10

    lrs_per_epoch_default = net.history[:, 'batches', :, 'default_lr']
    for lr in lrs_per_epoch_default:
        assert lr == [0.1] * 10

def test_tensorboard_logger(monkeypatch, data, tmpdir):
    add_scalar_mock = Mock()
    add_scalars_mock = Mock()

    monkeypatch.setattr('tensorboardX.SummaryWriter.add_scalar',
                        add_scalar_mock)
    monkeypatch.setattr('tensorboardX.SummaryWriter.add_scalars',
                        add_scalars_mock)

    log_dir = str(tmpdir.mkdir('tensorboardx_log'))

    module = MyModule()

    tensor_logger = TensorboardXLogger(
        log_dir,
        batch_targets=['train_loss'],
        epoch_targets=['valid_loss'],
        batch_groups=['train_loss'],
        epoch_groups=['valid_loss'])
    net = NeuralNet(
        module,
        torch.nn.MSELoss,
        callbacks=[tensor_logger],
        max_epochs=2,
        batch_size=100)

    X, y = data

    net.fit(X, y)

    batch_calls = [call('batch/train_loss', ANY, i) for i in range(8)]
    batch_calls += [call('batch/train_loss', ANY, i) for i in range(10, 18)]
    epoch_calls = [call('epoch/valid_loss', ANY, i) for i in range(1, 3)]
    add_scalar_mock.assert_has_calls(batch_calls + epoch_calls, any_order=True)

    batch_group_calls = [call('batch_train_loss', ANY, i) for i in range(8)]
    batch_group_calls += [
        call('batch_train_loss', ANY, i) for i in range(10, 18)
    ]
    epoch_group_calls = [call('epoch_valid_loss', ANY, i) for i in range(1, 3)]
    add_scalars_mock.assert_has_calls(
        batch_group_calls + epoch_group_calls, any_order=True)

max_stars_repo_name: thomasjpfan/mltome

max_issues_repo_name: thomasjpfan/mltome

max_forks_repo_name: thomasjpfan/mltome


text:
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

# Set random seed for reproducibility
np.random.seed(1000)

Nsamples = 10000

def X1_sample():
    return np.random.normal(0.1, 2.0)

def X2_sample(x1):
    return np.random.normal(x1, 0.5 + np.sqrt(np.abs(x1)))

if __name__ == '__main__':
    X = np.zeros((Nsamples,))
    Y = np.zeros((Nsamples,))

    for i, t in enumerate(range(Nsamples)):
        x1 = X1_sample()
        x2 = X2_sample(x1)

        X[i] = x1
        Y[i] = x2

    # Show the density estimation
    sns.set()

    fig, ax = plt.subplots(figsize=(10, 8))

    sns.kdeplot(X, Y, shade=True, shade_lowest=True, kernel="gau", ax=ax)
    ax.set_xlabel(r"$x_1$", fontsize=18)
    ax.set_ylabel(r"$x_2$", fontsize=18)
    ax.set_xlim([-5, 5])
    ax.set_ylim([-6, 6])

    plt.show()