190 lines
6.5 KiB
Python
190 lines
6.5 KiB
Python
#!/usr/bin env python3
|
|
import os
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import math
|
|
import scipy as scp
|
|
|
|
import matplotlib as mpl
|
|
mpl.rcParams["font.family"] = "serif"
|
|
mpl.rcParams["mathtext.fontset"] = "stix"
|
|
mpl.rcParams["text.usetex"] = True
|
|
mpl.rcParams['text.latex.preamble'] = f'''
|
|
\\usepackage{{amsmath}}
|
|
\\usepackage{{siunitx}}
|
|
\\input{{{os.path.abspath("../src/util/math-macros.tex")}}}
|
|
'''
|
|
|
|
|
|
if __name__ == "__main__": # make relative imports work as described here: https://peps.python.org/pep-0366/#proposed-change
|
|
if __package__ is None:
|
|
__package__ = "formulary"
|
|
import sys
|
|
filepath = os.path.realpath(os.path.abspath(__file__))
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(filepath)))
|
|
|
|
from util.mpl_colorscheme import set_mpl_colorscheme
|
|
import util.colorschemes as cs
|
|
from util.gen_tex_colorscheme import generate_latex_colorscheme
|
|
# SET THE COLORSCHEME
|
|
# hard white and black
|
|
# cs.p_gruvbox["fg0-hard"] = "#000000"
|
|
# cs.p_gruvbox["bg0-hard"] = "#ffffff"
|
|
COLORSCHEME = cs.gruvbox_light_no_beige()
|
|
# COLORSCHEME = cs.gruvbox_dark()
|
|
# cs.p_tum["fg0"] = cs.p_tum["alt-blue"]
|
|
# COLORSCHEME = cs.tum()
|
|
# COLORSCHEME = cs.legacy()
|
|
# COLORSCHEME = cs.stupid()
|
|
|
|
tex_aux_path = "../.aux/"
|
|
tex_src_path = "../src/"
|
|
img_out_dir = os.path.abspath(os.path.join(tex_src_path, "img"))
|
|
filetype = ".pdf"
|
|
skipasserts = False
|
|
|
|
def pt_2_inch(pt):
|
|
return 0.0138888889 * pt
|
|
def cm_2_inch(cm):
|
|
return 0.3937007874 * cm
|
|
|
|
# A4 - margins
|
|
width_line = cm_2_inch(21.0 - 2 * 2.0)
|
|
# width of a formula box, the prefactor has to match \eqwidth
|
|
width_formula = 0.69 * width_line
|
|
# arbitrary choice
|
|
height_default = width_line * 2 / 5
|
|
|
|
size_bigformula_fill_default = (width_line, height_default)
|
|
size_bigformula_half_quadratic = (width_line*0.5, width_line*0.5)
|
|
size_bigformula_small_quadratic = (width_line*0.33, width_line*0.33)
|
|
size_formula_fill_default = (width_formula, height_default)
|
|
size_formula_normal_default = (width_formula*0.8, height_default*0.8)
|
|
size_formula_half_quadratic = (width_formula*0.5, width_formula*0.5)
|
|
size_formula_small_quadratic = (width_formula*0.4, width_formula*0.4)
|
|
|
|
def assert_directory():
|
|
if not skipasserts:
|
|
assert os.path.abspath(".").endswith("scripts"), "Please run from the `scripts` directory"
|
|
|
|
def texvar(var, val, math=True):
|
|
s = "$" if math else ""
|
|
s += f"\\{var} = {val}"
|
|
if math: s += "$"
|
|
return s
|
|
|
|
def export(fig, name, tight_layout=True):
|
|
assert_directory()
|
|
filename = os.path.join(img_out_dir, name + filetype)
|
|
if tight_layout:
|
|
fig.tight_layout()
|
|
fig.savefig(filename, bbox_inches="tight", pad_inches=0.0)
|
|
|
|
|
|
def export_atoms(atoms, name, size, rotation="-30y,20x", get_bonds=True):
|
|
"""Export a render of ase atoms object"""
|
|
assert_directory()
|
|
wd = os.getcwd()
|
|
from util.aseutil import get_bondatoms, get_pov_settings
|
|
from ase import io
|
|
|
|
tmp_dir = os.path.join(os.path.abspath(tex_aux_path), "scripts_aux")
|
|
os.makedirs(tmp_dir, exist_ok=True)
|
|
os.chdir(tmp_dir)
|
|
|
|
out_filename = f"{name}.png"
|
|
|
|
bondatoms = None
|
|
if get_bonds:
|
|
bondatoms = get_bondatoms(atoms)
|
|
renderer = io.write(f'{name}.pov', atoms,
|
|
rotation=rotation,# text string with rotation (default='' )
|
|
radii=0.4, # float, or a list with one float per atom
|
|
show_unit_cell=2, # 0, 1, or 2 to not show, show, and show all of cell
|
|
colors=None, # List: one (r, g, b, t) tuple per atom
|
|
povray_settings=get_pov_settings(size, COLORSCHEME, bondatoms),
|
|
)
|
|
renderer.render()
|
|
os.chdir(wd)
|
|
os.rename(os.path.join(tmp_dir, out_filename), os.path.join(img_out_dir, out_filename))
|
|
|
|
@np.vectorize
|
|
def smooth_step(x: float, left_edge: float, right_edge: float):
|
|
x = (x - left_edge) / (right_edge - left_edge)
|
|
if x <= 0: return 0.
|
|
elif x >= 1: return 1.
|
|
else: return 3*(x*2) - 2*(x**3)
|
|
|
|
def blend_arrays(a_left, a_right, idx_left, idx_right, mode="linear"):
|
|
"""
|
|
Return a new array thats an overlap two arrays a_left and a_right.
|
|
Left of idx_left = a_left
|
|
Right of idx_right = a_right
|
|
Between: Smooth connection of both
|
|
"""
|
|
assert(a_left.shape == a_right.shape)
|
|
assert(idx_left < idx_right)
|
|
assert(idx_left > 0)
|
|
assert(idx_right < a_left.shape[0]-1)
|
|
ret = np.empty(a_left.shape)
|
|
ret[:idx_left] = a_left[:idx_left]
|
|
ret[idx_right:] = a_right[idx_right:]
|
|
n = idx_right - idx_left
|
|
left = a_left[idx_left]
|
|
right = a_right[idx_right]
|
|
for i in range(idx_left, idx_right):
|
|
j = i-idx_left # 0-based
|
|
if mode == "linear":
|
|
val = left * (n-j)/n + right * (j)/n
|
|
# connect with a single quadratic function
|
|
elif mode == "quadratic_simple":
|
|
slope_left = a_left[idx_left] - a_left[idx_left-1]
|
|
slope_right = a_right[idx_right+1] - a_right[idx_right]
|
|
b = slope_left
|
|
a = (slope_right-b)/(2*n)
|
|
c = left
|
|
# val = (left + slope_left * (j/n))*(j/n) + (right+slope_right*(1-j/n))*(1-j/n)
|
|
val = a*(j**2) + b*j +c
|
|
print(a,b,c)
|
|
# connect with two quadratic functions
|
|
# TODO: fix
|
|
elif mode == "quadratic":
|
|
slope_left = a_left[idx_left] - a_left[idx_left-1]
|
|
slope_right = a_right[idx_right+1] - a_right[idx_right]
|
|
c1 = left
|
|
b1 = slope_left
|
|
b2 = 2*slope_right - b1
|
|
a1 = (b2-b1)/4
|
|
a2 = -a1
|
|
c2 = right - a2*n**2-b2*n
|
|
m = 2 * (c2-c1)/(b1-b2)
|
|
print(m, a1, b1, c1)
|
|
print(m, a2, b2, c2)
|
|
|
|
if j < m:
|
|
val = a1*(j**2) + b1*j +c1
|
|
else:
|
|
val = a2*(j**2) + b2*j +c2
|
|
elif mode == "tanh":
|
|
TANH_FULL = 2 # x value where tanh is assumed to be 1/-1
|
|
x = (2 * j / n - 1) * TANH_FULL
|
|
tanh_error = 1-np.tanh(TANH_FULL)
|
|
amplitude = 0.5*(right-left)
|
|
val = amplitude * (1+tanh_error) * np.tanh(x) + (left+amplitude)
|
|
else: raise ValueError(f"Invalid mode: {mode}")
|
|
ret[i] = val
|
|
return ret
|
|
|
|
# run even when imported
|
|
set_mpl_colorscheme(COLORSCHEME)
|
|
|
|
if __name__ == "__main__":
|
|
assert_directory()
|
|
s = \
|
|
"""% This file was generated by scripts/formulary.py\n% Do not edit it directly, changes will be overwritten\n""" + generate_latex_colorscheme(COLORSCHEME)
|
|
filename = os.path.join(tex_src_path, "util/colorscheme.tex")
|
|
print(f"Writing tex colorscheme to {filename}")
|
|
with open(filename, "w") as file:
|
|
file.write(s)
|
|
|