pretty_print.py

import numpy as np

class CustomError(Exception):
       pass
    
def array_print(arrays,formats,return_text=False):
    '''function to print arrays in a nicer format for troubleshooting 
    and examples
    Arguments:
    ----------
    arrays = a list of arrays to be printed
    form = a list of formats for the arrays 
        each array should be surrounded by brackets
                    e.g. '[A]*', or '[x]', or, '=[b]'
    return_text (default) = False. if True, the text will be returned as an
                            output

    returns:
    ---------
    string = the printed string if needed again
    '''

    if len(arrays) != len(formats): 
        raise CustomError('the number of formats must be equal to the number of arrays')

    # create list of math operators on left-right of each array
    # math_op
    math_op=[[' ' for i in range(0,len(arrays[0]))] for j in range(0,2*len(arrays)+1)]
    math_loc=round(len(arrays[0])/2)-1
    for i,f in enumerate(formats):
        math_op[2*i][math_loc]=f.split('[')[0]+' '
        math_op[2*i+1][math_loc]=f.split(']')[-1]
    
    string=''
    for i in range(0,len(arrays[0])):
        for j,a in enumerate(arrays):
            a_row=np.atleast_1d(a[i])
            string+=('{}['+'{:7.2f}'*len(a_row)+']{}').format(
                math_op[2*j][i],*a_row,math_op[2*j+1][i])
        string+='\n'
    print(string)
    if return_text: return string

def array_latex(arrays,formats,return_text=False):
    '''function to print arrays in a latex format for viewing in 
    jupyter/markdown/latex equations

    Arguments:
    ----------
    arrays = a list of arrays to be printed
    form = a list of formats for the arrays 
        each array should be surrounded by brackets
                    e.g. '[A]*', or '[x]', or, '=[b]'
    return_text (default) = False. if True, the text will be returned as an
                            output

    returns:
    ---------
    string = the printed string if needed again
    '''

    if len(arrays) != len(formats): 
        raise CustomError('the number of formats must be equal to the number of arrays')

    # create list of math operators on left-right of each array
    # math_op
    math_op=[' ' for j in range(0,2*len(arrays)+1)]
    string=''
    for i,f in enumerate(formats):
      math_op[2*i]=f.split('[')[0]+' '
      math_op[2*i+1]=f.split(']')[-1]
    
    for j,a in enumerate(arrays):
      a_row=np.atleast_1d(a[0])
      string+=math_op[i]+'\\left[ \\begin{array}{'+'c'*len(a_row)+'}\n'
      i+=1
      for a_row in a:
          a_row=np.atleast_1d(a_row)
          string+=(' {:.1f} &'*len(a_row)).format(*a_row)[:-1]+'\\\\\n'
      string+='\\end{array}\\right]\n'+math_op[i]
      i+=1
    string+='\n'
    print(string)
    if return_text: return string