compat.py

import numpy as np
from mcerp import N, umath
from uncertainties import unumpy
import config as c



# ------------------------------------


def as_string(var):
    return f'{var=}'.split('=')[0]


def autotag(var):
    var.tag = as_string(var)


# ------------------------------------


# hack names
umath.arcsin = umath.asin 
unumpy.power = unumpy.pow

# unumpt funcs returns scalar array even as single output

COMPAT_LIBS = [np, unumpy, umath]


def try_op(name, *args, **kvargs):
    units = None
    if type(args[0]) == c.Q_:
        units = args[0].units

    for module in COMPAT_LIBS:
        try: 
            result = getattr(module, name)(*args, **kvargs)
            
            # convert scalar array to scalar
            if type(result) is np.ndarray and result.shape == ():
                result = result.item()

            # wrap as quantity if input was quantity
            if units != None and type(result) != c.Q_:
                result = c.Q_(result, units)

            return result
        except: continue
    raise Exception("Invalid op")



# ------------------------------------


def divide(a, b):
    return try_op("divide", a, b)

def power(a, b):
    return try_op("power", a, b)

def sum(a):
    return try_op("sum", a)

def log(a):
    return try_op("log", a)

def degrees(a):
    return try_op("degrees", a)

def radians(a):
    return try_op("radians", a)

def tan(a):
    return try_op("tan", a)

def cos(a):
    return try_op("cos", a)

def sin(a):
    return try_op("sin", a)

def arcsin(a):
    return try_op("arcsin", a)

def sqrt(a):
    return try_op("sqrt", a)