annotate yetilab/matrix/matrix.yeti @ 272:2ebda6646c40

Quicker, though uglier, sparse products

module yetilab.matrix.matrix;

// A matrix is an array of vectors.

// A matrix can be stored in either column-major (the default) or
// row-major format. Storage order is an efficiency concern only:
// every API function operating on matrix objects will return the same
// result regardless of storage order.  (The transpose function just
// switches the row/column order without moving the elements.)

vec = load yetilab.vector.vector;
bf = load yetilab.vector.blockfuncs;

load yetilab.vector.vectortype;
load yetilab.matrix.matrixtype;

size m =
    case m of
    DenseRows r:
        major = length r;
        { 
            rows = major, 
            columns = if major > 0 then vec.length r[0] else 0 fi,
        };
    DenseCols c:
        major = length c;
        { 
            rows = if major > 0 then vec.length c[0] else 0 fi,
            columns = major, 
        };
    SparseCSR { values, indices, pointers, extent }:
        {
            rows = (length pointers) - 1,
            columns = extent
        };
    SparseCSC { values, indices, pointers, extent }:
        {
            rows = extent,
            columns = (length pointers) - 1
        };
    esac;

width m = (size m).columns;
height m = (size m).rows;

nonZeroValues m =
   (nz d =
        sum
           (map do v:
                sum (map do n: if n == 0 then 0 else 1 fi done (vec.list v))
                done d);
    case m of 
    DenseRows d: nz d;
    DenseCols d: nz d;
    SparseCSR d: vec.length d.values;
    SparseCSC d: vec.length d.values;
    esac);

density m =
   ({ rows, columns } = size m;
    cells = rows * columns;
    (nonZeroValues m) / cells);

sparseSlice n d =
   (start = d.pointers[n];
    end = d.pointers[n+1];
    { 
        values = vec.slice d.values start end,
        indices = slice d.indices start end,
    });

nonEmptySlices d =
   (ne = array [];
    for [0..length d.pointers - 2] do i:
        if d.pointers[i] != d.pointers[i+1] then
            push ne i
        fi
    done;
    ne);

fromSlice n m d =
   (slice = sparseSlice n d;
    var v = 0;
    for [0..length slice.indices - 1] do i:
        if slice.indices[i] == m then
            v := vec.at slice.values i;
        fi
    done;
    v);

filledSlice n d =
   (slice = sparseSlice n d;
    dslice = new double[d.extent];
    for [0..length slice.indices - 1] do i:
        dslice[slice.indices[i]] := vec.at slice.values i;
    done;
    vec.vector dslice);

at' m row col =
    case m of
    DenseRows rows: r = rows[row]; vec.at r col;
    DenseCols cols: c = cols[col]; vec.at c row;
    SparseCSR data: fromSlice row col data;
    SparseCSC data: fromSlice col row data;
    esac;

getColumn j m =
    case m of
    DenseCols cols: cols[j];
    SparseCSC data: filledSlice j data;
    _: vec.fromList (map do i: at' m i j done [0..height m - 1]);
    esac;

getRow i m =
    case m of
    DenseRows rows: rows[i];
    SparseCSR data: filledSlice i data; 
    _: vec.fromList (map do j: at' m i j done [0..width m - 1]);
    esac;

asRows m =
    map do i: getRow i m done [0 .. (height m) - 1];

asColumns m =
    map do i: getColumn i m done [0 .. (width m) - 1];

isRowMajor? m =
    case m of
    DenseRows _: true;
    DenseCols _: false;
    SparseCSR _: true;
    SparseCSC _: false;
    esac;

isSparse? m =
    case m of
    DenseRows _: false;
    DenseCols _: false;
    SparseCSR _: true;
    SparseCSC _: true;
    esac;

typeOf m =
    if isRowMajor? m then RowMajor ()
    else ColumnMajor ()
    fi;

flippedTypeOf m =
    if isRowMajor? m then ColumnMajor ()
    else RowMajor ()
    fi;

newColumnMajorStorage { rows, columns } = 
    if rows < 1 then array []
    else array (map \(vec.zeros rows) [1..columns])
    fi;

zeroMatrix { rows, columns } = 
    DenseCols (newColumnMajorStorage { rows, columns });

zeroMatrixWithTypeOf m { rows, columns } = 
    if isRowMajor? m then
        DenseRows (newColumnMajorStorage { rows = columns, columns = rows });
    else
        DenseCols (newColumnMajorStorage { rows, columns });
    fi;

zeroSizeMatrix () = zeroMatrix { rows = 0, columns = 0 };

generate f { rows, columns } =
    if rows < 1 or columns < 1 then zeroSizeMatrix ()
    else
        m = array (map \(new double[rows]) [1..columns]);
        for [0..columns-1] do col:
            for [0..rows-1] do row:
                m[col][row] := f row col;
            done;
        done;
        DenseCols (array (map vec.vector m))
    fi;

swapij =
    map do { i, j, v }: { i = j, j = i, v } done;

//!!! should use { row = , column = , value = } instead of i, j, v?
enumerateSparse m =
   (enumerate { values, indices, pointers } =
        concat
           (map do i:
                start = pointers[i];
                end = pointers[i+1];
                map2 do j v: { i, j, v } done 
                    (slice indices start end)
                    (vec.list (vec.slice values start end))
                done [0..length pointers - 2]);
    case m of
    SparseCSC d: swapij (enumerate d);
    SparseCSR d: enumerate d;
     _: [];
    esac);

enumerateDense m =
   (enumerate d =
        concat
           (map do i:
                vv = d[i];
                map2 do j v: { i, j, v } done
                    [0..vec.length vv - 1]
                    (vec.list vv);
                done [0..length d - 1]);
    case m of
    DenseCols c: swapij (enumerate c);
    DenseRows r: enumerate r;
     _: [];
    esac);

enumerate m =
    if isSparse? m then enumerateSparse m else enumerateDense m fi;

// Make a sparse matrix from entries whose i, j values are known to be
// within range
makeSparse type size data =
   (isRow = case type of RowMajor (): true; ColumnMajor (): false esac;
    ordered = 
        sortBy do a b:
            if a.maj == b.maj then a.min < b.min else a.maj < b.maj fi
        done
           (map
                if isRow then
                    do { i, j, v }: { maj = i, min = j, v } done;
                else
                    do { i, j, v }: { maj = j, min = i, v } done;
                fi
               (filter do d: d.v != 0 done data));
    tagger = if isRow then SparseCSR else SparseCSC fi;
    majorSize = if isRow then size.rows else size.columns fi;
    minorSize = if isRow then size.columns else size.rows fi;
    pointers = array [0];
    setArrayCapacity pointers (size.rows + 1);
    fillPointers n i data =
        if n < majorSize then
            case data of
            d::rest:
               (for [n..d-1] \(push pointers i);
                fillPointers d (i+1) rest);
             _:
                for [n..majorSize-1] \(push pointers i);
            esac;
        fi;
    fillPointers 0 0 (map (.maj) ordered);
    tagger {
        values = vec.fromList (map (.v) ordered),
        indices = array (map (.min) ordered),
        pointers,
        extent = minorSize,
    });

// Make a sparse matrix from entries that may contain out-of-range
// cells which need to be filtered out. This is the public API for
// makeSparse and is also used to discard out-of-range cells from
// resizedTo.
newSparseMatrix type size data =
    makeSparse type size
       (filter
            do { i, j, v }:
                i == int i and i >= 0 and i < size.rows and 
                j == int j and j >= 0 and j < size.columns
            done data);

toSparse m =
    if isSparse? m then m
    else
        makeSparse (typeOf m) (size m) (enumerateDense m);
    fi;

toDense m =
    if not (isSparse? m) then m
    elif isRowMajor? m then
        DenseRows (array (map do row: getRow row m done [0..height m - 1]));
    else
        DenseCols (array (map do col: getColumn col m done [0..width m - 1]));
    fi;

constMatrix n = generate do row col: n done;
randomMatrix = generate do row col: Math#random() done;
identityMatrix = constMatrix 1;

transposed m =
    case m of
    DenseRows d: DenseCols d;
    DenseCols d: DenseRows d;
    SparseCSR d: SparseCSC d;
    SparseCSC d: SparseCSR d;
    esac;

flipped m =
    if isSparse? m then
        makeSparse (flippedTypeOf m) (size m) (enumerateSparse m)
    else
        if isRowMajor? m then
            generate do row col: at' m row col done (size m);
        else
            transposed
               (generate do row col: at' m col row done
                { rows = (width m), columns = (height m) });
        fi
    fi;

toRowMajor m =
    if isRowMajor? m then m else flipped m fi;

toColumnMajor m =
    if not isRowMajor? m then m else flipped m fi;

equal'' comparator vecComparator m1 m2 =
    // Prerequisite: m1 and m2 have same sparse-p and storage order
   (compareVecLists vv1 vv2 = all id (map2 vecComparator vv1 vv2);
    compareSparse d1 d2 =
        d1.extent == d2.extent and
        vecComparator d1.values d2.values and
        d1.indices == d2.indices and
        d1.pointers == d2.pointers;
    case m1 of
    DenseRows d1:
        case m2 of DenseRows d2: compareVecLists d1 d2; _: false; esac;
    DenseCols d1:
        case m2 of DenseCols d2: compareVecLists d1 d2; _: false; esac;
    SparseCSR d1:
        case m2 of SparseCSR d2: compareSparse d1 d2; _: false; esac;
    SparseCSC d1:
        case m2 of SparseCSC d2: compareSparse d1 d2; _: false; esac;
    esac);

equal' comparator vecComparator m1 m2 =
    if size m1 != size m2 then 
        false
    elif isRowMajor? m1 != isRowMajor? m2 then
        equal' comparator vecComparator (flipped m1) m2;
    elif isSparse? m1 != isSparse? m2 then
        if isSparse? m1 then
            equal' comparator vecComparator m1 (toSparse m2)
        else
            equal' comparator vecComparator (toSparse m1) m2
        fi
    else
        equal'' comparator vecComparator m1 m2
    fi;

// Compare matrices using the given comparator for individual cells.
// Note that matrices with different storage order but the same
// contents are equal, although comparing them is slow.
//!!! Document the fact that sparse matrices can only be equal if they
// have the same set of non-zero cells (regardless of comparator used)
equalUnder comparator =
    equal' comparator (vec.equalUnder comparator);

equal =
    equal' (==) vec.equal;

newMatrix type data = //!!! NB does not copy data
   (tagger = case type of RowMajor (): DenseRows; ColumnMajor (): DenseCols esac;
    if empty? data or vec.empty? (head data)
    then zeroSizeMatrix ()
    else tagger (array data)
    fi);

newRowVector data = //!!! NB does not copy data
    DenseRows (array [data]);

newColumnVector data = //!!! NB does not copy data
    DenseCols (array [data]);

denseLinearOp op m1 m2 =
    if isRowMajor? m1 then
        newMatrix (typeOf m1) 
           (map2 do c1 c2: op c1 c2 done (asRows m1) (asRows m2));
    else
        newMatrix (typeOf m1) 
           (map2 do c1 c2: op c1 c2 done (asColumns m1) (asColumns m2));
    fi;

sparseSumOrDifference op m1 m2 =
   (h = [:];
    for (enumerate m1) do { i, j, v }:
        if not (i in h) then h[i] := [:] fi;
        h[i][j] := v;
    done;
    for (enumerate m2) do { i, j, v }:
        if not (i in h) then h[i] := [:] fi;
        if j in h[i] then h[i][j] := op h[i][j] v;
        else h[i][j] := op 0 v;
        fi;
    done;
    entries = concat
       (map do i:
            kk = keys h[i];
            map2 do j v: { i, j, v } done kk (map (at h[i]) kk)
            done (keys h));
    makeSparse (typeOf m1) (size m1) entries);

sum' m1 m2 =
    if (size m1) != (size m2)
    then failWith "Matrices are not the same size: \(size m1), \(size m2)";
    elif isSparse? m1 and isSparse? m2 then
        sparseSumOrDifference (+) m1 m2;
    else
        denseLinearOp bf.add m1 m2;
    fi;

difference m1 m2 =
    if (size m1) != (size m2)
    then failWith "Matrices are not the same size: \(size m1), \(size m2)";
    elif isSparse? m1 and isSparse? m2 then
        sparseSumOrDifference (-) m1 m2;
    else
        denseLinearOp bf.subtract m1 m2;
    fi;

scaled factor m =
    if isSparse? m then
        makeSparse (typeOf m) (size m)
           (map do { i, j, v }: { i, j, v = factor * v } done (enumerate m))
    elif isRowMajor? m then
        newMatrix (typeOf m) (map (bf.scaled factor) (asRows m));
    else
        newMatrix (typeOf m) (map (bf.scaled factor) (asColumns m));
    fi;

abs' m =
    if isSparse? m then
        makeSparse (typeOf m) (size m)
           (map do { i, j, v }: { i, j, v = abs v } done (enumerate m))
    elif isRowMajor? m then
        newMatrix (typeOf m) (map bf.abs (asRows m));
    else
        newMatrix (typeOf m) (map bf.abs (asColumns m));
    fi;

filter f m =
    if isSparse? m then
        makeSparse (typeOf m) (size m)
           (map do { i, j, v }: { i, j, v = if f v then v else 0 fi } done
               (enumerate m))
    else
        vfilter = vec.fromList . (map do i: if f i then i else 0 fi done) . vec.list;
        if isRowMajor? m then
            newMatrix (typeOf m) (map vfilter (asRows m));
        else
            newMatrix (typeOf m) (map vfilter (asColumns m));
        fi;
    fi;

sparseProductLeft size m1 m2 =
   ({ values, indices, pointers } = case m1 of
         SparseCSR d: d;
         SparseCSC d: d;
         _: failWith "sparseProductLeft called for non-sparse m1";
         esac;
    rows = isRowMajor? m1;
    data = array (map \(new double[size.rows]) [1..size.columns]);
    for [0..size.columns - 1] do j':
        c = getColumn j' m2;
        var p = 0;
        for [0..length indices - 1] do ix:
            ix == pointers[p+1] loop (p := p + 1);
            i = if rows then p else indices[ix] fi;
            j = if rows then indices[ix] else p fi;
            data[j'][i] := data[j'][i] + (vec.at values ix) * (vec.at c j);
        done;
    done;
    DenseCols (array (map vec.vector (list data))));

sparseProductRight size m1 m2 =
   ({ values, indices, pointers } = case m2 of
         SparseCSR d: d;
         SparseCSC d: d;
         _: failWith "sparseProductLeft called for non-sparse m1";
         esac;
    rows = isRowMajor? m2;
    data = array (map \(new double[size.columns]) [1..size.rows]);
    for [0..size.rows - 1] do i':
        r = getRow i' m1;
        var p = 0;
        for [0..length indices - 1] do ix:
            ix == pointers[p+1] loop (p := p + 1);
            i = if rows then p else indices[ix] fi;
            j = if rows then indices[ix] else p fi;
            data[i'][j] := data[i'][j] + (vec.at values ix) * (vec.at r i);
        done;
    done;
    DenseRows (array (map vec.vector (list data))));

sparseProduct size m1 m2 =
    case m2 of
    SparseCSC d:
       ({ values, indices, pointers } = case m1 of
            SparseCSR d1: d1;
            SparseCSC d1: d1;
            _: failWith "sparseProduct called for non-sparse matrices";
            esac;
        rows = isRowMajor? m1;
        var p = 0;
        pindices = new int[length indices];
        for [0..length indices - 1] do ix:
            ix == pointers[p+1] loop (p := p + 1);
            pindices[ix] := p;
        done;
        entries =
           (map do j':
                cs = sparseSlice j' d;
                hin = mapIntoHash
                   (at cs.indices) (vec.at cs.values)
                   [0..length cs.indices - 1];
                hout = [:];
                for [0..length indices - 1] do ix:
                    i = if rows then pindices[ix] else indices[ix] fi;
                    j = if rows then indices[ix] else pindices[ix] fi;
                    if j in hin then
                        p = (vec.at values ix) * hin[j];
                        hout[i] := p + (if i in hout then hout[i] else 0 fi);
                    fi;
                done;
                map do i:
                    { i, j = j', v = hout[i] }
                done (keys hout);
            done (nonEmptySlices d));
        makeSparse (ColumnMajor ()) size (concat entries));
    SparseCSR _:
        sparseProduct size m1 (flipped m2);
     _: failWith "sparseProduct called for non-sparse matrices";
    esac;

denseProduct size m1 m2 =
   (data = array (map \(new double[size.rows]) [1..size.columns]);
    for [0..size.rows - 1] do i:
        row = getRow i m1;
        for [0..size.columns - 1] do j:
            data[j][i] := bf.sum (bf.multiply row (getColumn j m2));
        done;
    done;
    DenseCols (array (map vec.vector (list data))));

product m1 m2 =
    if (size m1).columns != (size m2).rows
    then failWith "Matrix dimensions incompatible: \(size m1), \(size m2) (\((size m1).columns) != \((size m2).rows))";
    else 
        size = { rows = (size m1).rows, columns = (size m2).columns };
        if isSparse? m1 then
            if isSparse? m2 then
                sparseProduct size m1 m2
            else
                sparseProductLeft size m1 m2
            fi
        elif isSparse? m2 then
            sparseProductRight size m1 m2
        else
            denseProduct size m1 m2
        fi;
    fi;

concatAgainstGrain tagger getter counter mm =
   (n = counter (size (head mm));
    tagger (array
       (map do i:
           vec.concat (map (getter i) mm)
           done [0..n-1])));

concatWithGrain tagger getter counter mm =
    tagger (array
       (concat
           (map do m:
               n = counter (size m);
               map do i: getter i m done [0..n-1]
               done mm)));

sparseConcat direction first mm =
   (dimension d f = if direction == d then sum (map f mm) else f first fi;
    rows = dimension (Vertical ()) height;
    columns = dimension (Horizontal ()) width;
    entries ioff joff ui uj mm =
        case mm of 
        m::rest:
           (map do { i, j, v }: { i = i + ioff, j = j + joff, v }
                done (enumerate m)) ++
           (entries
               (ioff + ui * height m)
               (joff + uj * width m)
                ui uj rest);
         _: []
        esac;
    makeSparse (typeOf first) { rows, columns }
        if direction == Vertical () then entries 0 0 1 0 mm
        else entries 0 0 0 1 mm fi);

checkDimensionsFor direction first mm =
   (counter = if direction == Horizontal () then (.rows) else (.columns) fi;
    n = counter (size first);
    if not (all id (map do m: counter (size m) == n done mm)) then
        failWith "Matrix dimensions incompatible for concat (found \(map do m: counter (size m) done mm) not all of which are \(n))";
    fi);

concat direction mm = //!!! doc: storage order is taken from first matrix in sequence
    case length mm of
    0: zeroSizeMatrix ();
    1: head mm;
    _:
        first = head mm;
        checkDimensionsFor direction first mm;
        if all isSparse? mm then
            sparseConcat direction first mm
        else
            row = isRowMajor? first;
            // horizontal, row-major: against grain with rows
            // horizontal, col-major: with grain with cols
            // vertical, row-major: with grain with rows
            // vertical, col-major: against grain with cols
            case direction of
            Horizontal ():
                if row then concatAgainstGrain DenseRows getRow (.rows) mm;
                else concatWithGrain DenseCols getColumn (.columns) mm;
                fi;
            Vertical ():
                if row then concatWithGrain DenseRows getRow (.rows) mm;
                else concatAgainstGrain DenseCols getColumn (.columns) mm;
                fi;
            esac;
        fi;
    esac;

//!!! doc note: argument order chosen for consistency with std module slice
rowSlice m start end = //!!! doc: storage order same as input
    if isRowMajor? m then
        DenseRows (array (map ((flip getRow) m) [start .. end - 1]))
    else 
        DenseCols (array (map do v: vec.slice v start end done (asColumns m)))
    fi;

//!!! doc note: argument order chosen for consistency with std module slice
columnSlice m start end = //!!! doc: storage order same as input
    if not isRowMajor? m then
        DenseCols (array (map ((flip getColumn) m) [start .. end - 1]))
    else 
        DenseRows (array (map do v: vec.slice v start end done (asRows m)))
    fi;

resizedTo newsize m =
   (if newsize == (size m) then
        m
    elif isSparse? m then
        // don't call makeSparse directly: want to discard
        // out-of-range cells
        newSparseMatrix (typeOf m) newsize (enumerateSparse m)
    elif (height m) == 0 or (width m) == 0 then
        zeroMatrixWithTypeOf m newsize;
    else
        growrows = newsize.rows - (height m);
        growcols = newsize.columns - (width m);
        rowm = isRowMajor? m;
        resizedTo newsize
            if rowm and growrows < 0 then
                rowSlice m 0 newsize.rows
            elif (not rowm) and growcols < 0 then 
                columnSlice m 0 newsize.columns
            elif growrows < 0 then 
                rowSlice m 0 newsize.rows
            elif growcols < 0 then 
                columnSlice m 0 newsize.columns
            else
                if growrows > 0 then
                    concat (Vertical ())
                       [m, zeroMatrixWithTypeOf m ((size m) with { rows = growrows })]
                else
                    concat (Horizontal ())
                       [m, zeroMatrixWithTypeOf m ((size m) with { columns = growcols })]
                fi
            fi
    fi);

{
    size,
    width,
    height,
    density,
    nonZeroValues,
    at = at',
    getColumn,
    getRow,
    isRowMajor?,
    isSparse?,
    generate,
    constMatrix,
    randomMatrix,
    zeroMatrix,
    identityMatrix,
    zeroSizeMatrix,
    equal,
    equalUnder,
    transposed,
    flipped,
    toRowMajor,
    toColumnMajor,
    toSparse,
    toDense,
    scaled,
    resizedTo,
    asRows,
    asColumns,
    sum = sum',
    difference,
    abs = abs',
    filter,
    product,
    concat,
    rowSlice,
    columnSlice,
    newMatrix,
    newRowVector,
    newColumnVector,
    newSparseMatrix,
    enumerate
}
as
{
//!!! check whether these are right to be .selector rather than just selector

    size is matrix -> { .rows is number, .columns is number },
    width is matrix -> number,
    height is matrix -> number,
    density is matrix -> number,
    nonZeroValues is matrix -> number,
    at is matrix -> number -> number -> number,
    getColumn is number -> matrix -> vector,
    getRow is number -> matrix -> vector,
    isRowMajor? is matrix -> boolean,
    isSparse? is matrix -> boolean,
    generate is (number -> number -> number) -> { .rows is number, .columns is number } -> matrix,
    constMatrix is number -> { .rows is number, .columns is number } -> matrix,
    randomMatrix is { .rows is number, .columns is number } -> matrix,
    zeroMatrix is { .rows is number, .columns is number } -> matrix, 
    identityMatrix is { .rows is number, .columns is number } -> matrix, 
    zeroSizeMatrix is () -> matrix,
    equal is matrix -> matrix -> boolean,
    equalUnder is (number -> number -> boolean) -> matrix -> matrix -> boolean,
    transposed is matrix -> matrix,
    flipped is matrix -> matrix, 
    toRowMajor is matrix -> matrix, 
    toColumnMajor is matrix -> matrix,
    toSparse is matrix -> matrix,
    toDense is matrix -> matrix,
    scaled is number -> matrix -> matrix,
    thresholded is number -> matrix -> matrix,
    resizedTo is { .rows is number, .columns is number } -> matrix -> matrix,
    asRows is matrix -> list<vector>, 
    asColumns is matrix -> list<vector>,
    sum is matrix -> matrix -> matrix,
    difference is matrix -> matrix -> matrix,
    abs is matrix -> matrix,
    filter is (number -> boolean) -> matrix -> matrix,
    product is matrix -> matrix -> matrix,
    concat is (Horizontal () | Vertical ()) -> list<matrix> -> matrix,
    rowSlice is matrix -> number -> number -> matrix, 
    columnSlice is matrix -> number -> number -> matrix,
    newMatrix is (ColumnMajor () | RowMajor ()) -> list<vector> -> matrix, 
    newRowVector is vector -> matrix, 
    newColumnVector is vector -> matrix,
    newSparseMatrix is (ColumnMajor () | RowMajor ()) -> { .rows is number, .columns is number } -> list<{ .i is number, .j is number, .v is number }> -> matrix,
    enumerate is matrix -> list<{ .i is number, .j is number, .v is number }>
}