formelsammlung/scripts/formulary.py

#!/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'] = r'\usepackage{amsmath}\usepackage{siunitx}'


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()
# 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)


# 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)