small_functions.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Module for tiny helper functions, such as path-string manipulation.

small helpful functions

FUNCTIONS:
    is_nondigit(path): False if path tail contains digits
    timestep_in(path): true if 'timestep' is in path tail
    check_simdir(path): param dict based on valid path (falsy {} if invalid)
    a_in_b_dict(a, b): True if a[key] in b[key] for all b.keys

Created on Thu Aug 12 14:28:37 2021

@author: yoav
"""

import os
import re

import numpy as np


def is_nondigit(path) -> bool:
    """Return True if all characters in path tail are not digit."""
    return not any(x.isdigit() for x in os.path.split(path)[-1])


def timestep_in(path) -> bool:
    """Return True if path tail contains "timestep"."""
    return 'timestep' in os.path.split(path)[-1]


_pattern_vsd=re.compile('[^\W\d_]+|\d+')  # from https://stackoverflow.com/questions/12409894
def valid_simulation_dict(path) -> dict:
    """Check path as a valid simulation directory, return {param: index}.

    Returns a dictionary of {parameter: index} if valid, else a falsy {}.
    A directory is a valid simulation directory if it is of the form
    ab12c3... producing the dictionary {'ab': 12, 'c': 3, ...}
    or of the form x####_ab12c3... 
    where the x#### is the simulation number which is skipped
    """
    # validate parameter section by constructing parameter dictionary
    # split to numeric-alpha words ['f','102','blah',987]
    wordlist=_pattern_vsd.findall(os.path.split(os.path.realpath(path))[-1])
    if wordlist and wordlist[0]=='x':
        # Skip initial x:num (simulation number)
        wordlist = wordlist[2:]
    try:
        params = {key: int(value_word) for key, value_word in
                  zip(wordlist[::2], wordlist[1::2])}
    except (ValueError, TypeError):
        # unable to construct dictionary
        return {}
    return params


def compile_tape_regex_option(option_list):
    """Return a regex pattern that can extract the option in option_list from a tape file using pattern.findall(tape)"""
    if option_list is None or option_list is all:
        return re.compile("^([\w\d]*)=(.*)",re.MULTILINE)
    if type(option_list) is str:
        re.compile(f"^({option_list})=(.*)",re.MULTILINE)
    return re.compile(f"^({'|'.join(map(str,option_list))})=(.*)",re.MULTILINE)


def _string_to_value(value):
    """Conert string to value, int first, then double"""
    try:
        return int(value)
    except ValueError:
        try:
            return float(value)
        except ValueError:
            return value

_pattern_tape_all=compile_tape_regex_option(None)
def get_tape_options(tape):
    """Extract tape as text to a dictionary of all options"""
    options = _pattern_tape_all.findall(tape)
    return {option: _string_to_value(value) for option,value in options}

_pattern_tape_adhesion=compile_tape_regex_option(['z_adhesion','adhesion_z','adhesion_cuttoff','adhesion_cutoff'])
def extract_tape_adhesion_plane(tape):
    """Get the top of the adhesion plane from tape (z0+cutoff).
    
    This means vertex position below this value are in adhesion.
    """
    z_adh, z_cutoff = _pattern_tape_adhesion.findall(tape)
    return float(z_adh[1])+float(z_cutoff[1])

_pattern_tape_adhesion_2=compile_tape_regex_option(['z_adhesion','adhesion_z','adhesion_cuttoff','adhesion_cutoff','adhesion_radius'])
def from_tape_is_adhered_to_sphere(tape, pos):
    """Get the top of the adhesion sphere from tape (r+cutoff).
    
    Returns boolean array from pos
    """
    cutoff, radius, z0 = sorted(_pattern_tape_adhesion_2.findall(tape), key=lambda x: x[0])
    return ((pos-[[0,0,float(z0[1])-float(radius[1])]])**2).sum(axis=-1)<(float(cutoff[1])+float(radius[1]))**2


def a_in_b_dict(a: dict, b: dict) -> bool:
    """Return True if for all keys a[key] is in b[key]."""
    return all(a[key] in b[key] for key in b.keys())


def np_index(arr: np.ndarray, val):
    """Numpy array equivalent to a list's [].index method."""
    return np.where(arr == val)[0][0]


def subspace_indices(sub_space, param_idx, keys):
    """Get indice of {'a': 2} in {'a': [1,2,3]}."""
    return [np_index(sub_space[key], param_idx[key]) for key in keys]


def str_to_slice(string: str) -> slice:
    """Convert string representation to slice."""
    if ':' not in string:
        x = int(string)
        if x == -1:
            return slice(x, None)
        return slice(x, x+1)
    return slice(*map(lambda x: int(x.strip()) if x.strip() else None,
                      string.split(':')))


def _int_slice_to_slice(s):
    """Convert to slice: internal function for getitem."""
    if type(s) is int:
        if s == -1:
            return slice(-1, None)
        return slice(s, s+1)
    return s


class _SliceGetter():
    """Slice()[::2] -> slice(None,None,2)."""

    def __getitem__(self, s):
        """Return slice generated by [::]."""
        if type(s) is tuple:
            return tuple(_int_slice_to_slice(x) for x in s)
        return _int_slice_to_slice(s)


Slice = _SliceGetter()  # global to export: Slice[::2]


class _RangeGetter():
    """Range(5)[::2] -> range(5)[::2]."""

    def __init__(self, N=None):
        self.N = N

    def slice_to_range(self, s):
        if self.N is not None:
            return range(self.N)[s]
        else:
            a, b, c = s.start, s.step, s.stop
            b = b if b is not None else 1
            a = a if a is not None else 0 if b > 0 else -1
            c = c if c is not None else a+b+1 if b is not None else a+1
            return range(a, c, b)

    def __getitem__(self, s):
        """Return slice generated by [::]."""
        if type(s) is tuple:
            return tuple(self.slice_to_range(_int_slice_to_slice(x)) for x in s)
        return self.slice_to_range(_int_slice_to_slice(s))

    def __call__(self, N):
        return _RangeGetter(N)


Range = _RangeGetter()


def len_slice_max(s: slice) -> int:
    """Get max length of slice.

    If possibly inf, return None.
    """
    start = s.start if s.start is not None else 0
    step = s.step if s.step is not None else 1
    stop = s.stop
    if start < 0:
        if stop is None:
            return (-start)//step
        elif stop > 0:
            return max((-start)//step, stop//step)
        elif stop < 0:
            if stop > start:
                return (stop-start)//step
            else:
                return 0
    else:  # start>=0
        if stop is not None and stop > start:
            return (stop-start)//step
    return None


def upmask(mask1, mask2):
    """Given A[mask1]=B and B[mask2]=c, return mask3 such that A[mask3]=c"""
    m = mask1.copy()
    m[mask1] = mask2
    return m


def normalize_axis(vectors, axis=-1):
    """Take an array, normalize the given axis e.g. A/norm(A,axis)"""
    n = np.linalg.norm(vectors, axis=axis, keepdims=True)
    n[n == 0] = 1
    return vectors/n