tri.hpp

Go to the documentation of this file.

#pragma once
 
#include "NumCpp/Core/Internal/StaticAsserts.hpp"
#include "NumCpp/Core/Shape.hpp"
#include "NumCpp/Core/Types.hpp"
#include "NumCpp/NdArray.hpp"
 
namespace nc
{
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> tril(uint32 inN, int32 inOffset = 0)
    {
        STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
        uint32 rowStart = 0;
        uint32 colStart = 0;
        if (inOffset > 0)
        {
            colStart = inOffset;
        }
        else
        {
            rowStart = inOffset * -1;
        }
 
        NdArray<dtype> returnArray(inN);
        returnArray.zeros();
        for (uint32 row = rowStart; row < inN; ++row)
        {
            for (uint32 col = 0; col < row + colStart + 1 - rowStart; ++col)
            {
                if (col == inN)
                {
                    break;
                }
 
                returnArray(row, col) = dtype{ 1 };
            }
        }
 
        return returnArray;
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> tril(uint32 inN, uint32 inM, int32 inOffset = 0)
    {
        STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
        uint32 rowStart = 0;
        uint32 colStart = 0;
        if (inOffset > 0)
        {
            colStart = inOffset;
        }
        else if (inOffset < 0)
        {
            rowStart = inOffset * -1;
        }
 
        NdArray<dtype> returnArray(inN, inM);
        returnArray.zeros();
        for (uint32 row = rowStart; row < inN; ++row)
        {
            for (uint32 col = 0; col < row + colStart + 1 - rowStart; ++col)
            {
                if (col == inM)
                {
                    break;
                }
 
                returnArray(row, col) = dtype{ 1 };
            }
        }
 
        return returnArray;
    }
 
    // forward declare
    template<typename dtype>
    NdArray<dtype> triu(uint32 inN, uint32 inM, int32 inOffset = 0);
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> tril(const NdArray<dtype>& inArray, int32 inOffset = 0)
    {
        STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
        const Shape inShape  = inArray.shape();
        auto        outArray = inArray.copy();
        outArray.putMask(triu<bool>(inShape.rows, inShape.cols, inOffset + 1), 0);
        return outArray;
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> triu(uint32 inN, uint32 inM, int32 inOffset)
    {
        STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
        // because i'm stealing the lines of code from tril and reversing it, this is necessary
        inOffset -= 1;
 
        uint32 rowStart = 0;
        uint32 colStart = 0;
        if (inOffset > 0)
        {
            colStart = inOffset;
        }
        else if (inOffset < 0)
        {
            rowStart = inOffset * -1;
        }
 
        NdArray<dtype> returnArray(inN, inM);
        returnArray.ones();
        for (uint32 row = rowStart; row < inN; ++row)
        {
            for (uint32 col = 0; col < row + colStart + 1 - rowStart; ++col)
            {
                if (col == inM)
                {
                    break;
                }
 
                returnArray(row, col) = dtype{ 0 };
            }
        }
 
        return returnArray;
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> triu(uint32 inN, int32 inOffset = 0)
    {
        STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
        return tril<dtype>(inN, -inOffset).transpose();
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<dtype> triu(const NdArray<dtype>& inArray, int32 inOffset = 0)
    {
        STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
        const Shape inShape  = inArray.shape();
        auto        outArray = inArray.copy();
        outArray.putMask(tril<bool>(inShape.rows, inShape.cols, inOffset - 1), 0);
        return outArray;
    }
} // namespace nc