<filename>FADiff/FuncVect.py
#!/usr/bin/env python3

from FADiff import FADiff
from fad.Gradients import Scal as fadScal
from fad.Matrices import Vect as fadVect
from rev.Gradients import Scal as revScal
from rev.Matrices import Vect as revVect
import numpy as np


class FuncVect:
    """
    This class represents a vector function for automatic differention (AD) for
    forward and reverse mode.
    """
    def __init__(self, funcs):
        """
        Inputs:
            funcs : int/float
                list of functions where all functions in the list are either
                all Scal objects or all Vect objects and all reverse mode or all
                forward mode objects
        """
        try:                                   # All fxns all Scal xor all Vect?
            func_type = type(funcs[0])         # Check the first element
            for func in funcs[1:]:
                if type(func) != func_type:    # Check for consistency
                    raise Exception('All functions must be of same type (Scalar '
                                    'or Vector).')
        except Exception:
            raise
        else:
            # Set inputs list based on type (Scal or Vect) of functions
            self._inputs = []
            if func_type is fadScal:
                self._inputs = FADiff._fadscal_inputs
            elif func_type is fadVect:
                self._inputs = FADiff._fadvect_inputs
            elif func_type is revScal:
                self._inputs = FADiff._revscal_inputs
            elif func_type is revVect:
                self._inputs = FADiff._revvect_inputs
            # List of objects (Scal or Vect) used in fxns
            self._f_vect = funcs
            self._input_vars = []  # Get complete list of input vars of f_vect
            for func in funcs:
                try:
                    if func._parents:  # Has parents?
                        for var in func._der.keys():  # Loop through existing roots
                            if var in func._parents:  # If var is in parent/root list
                                self._input_vars.append(var)
                    # Otherwise input var, i.e., f(x) = x?
                    elif func in self._inputs:
                        self._input_vars.append(func)
                except AttributeError:
                    for var in func._inputs.keys():
                        self._input_vars.append(var)
            self._input_vars = list(set(self._input_vars))

    @property
    def val(self):
        """
        Returns NumPy array of values
        """
        func_vals = []
        for func in self._f_vect:   # For each fxn, add its value to list
            func_vals.append(func._val)
        return np.squeeze(np.array(func_vals))  # Return list

    @property
    def der(self):
        """
        Returns partial derivatives of function vector wrt all root input vars
        used in all functions in the vector in a numpy array
        """
        all_fxns_jacobs = []
        for func in self._f_vect:
            one_fxn_jacob = []
            try:
                for var, part_der in func._der.items():
                    if var in self._input_vars:
                        one_fxn_jacob.append(part_der)
            except AttributeError:
                func_der = func.der
                idx = 0
                for var in self._inputs:  # Using this keeps order of vars
                    if var in self._input_vars:
                        if var in func._inputs.keys():
                            one_fxn_jacob.append(func_der[idx])
                            idx += 1
                        else:
                            one_fxn_jacob.append(0)
            all_fxns_jacobs.append(one_fxn_jacob)
        return np.squeeze(np.array(all_fxns_jacobs))
