updating tests

updating tests

use NumSuch,
    Norm,
    LinearAlgebra,
    Viterbi,
    Core,
    Charcoal;

class NumSuchTest : UnitTest {
  var nv: int = 8,
      D: domain(2) = {1..nv, 1..nv},
      SD: sparse subdomain(D),
      X: [SD] real;

  var vn: [1..0] string;

  proc setUp() {
    vn.push_back("star lord");
    vn.push_back("gamora");
    vn.push_back("groot");
    vn.push_back("drax");
    vn.push_back("rocket");
    vn.push_back("mantis");
    vn.push_back("yondu");
    vn.push_back("nebula");

    SD += (1,2); X[1,2] = 1;
    SD += (1,3); X[1,3] = 1;
    SD += (1,4); X[1,4] = 1;
    SD += (2,2); X[2,2] = 1;
    SD += (2,4); X[2,4] = 1;
    SD += (3,4); X[3,4] = 1;
    SD += (4,5); X[4,5] = 1;
    SD += (5,6); X[5,6] = 1;
    SD += (6,7); X[6,7] = 1;
    SD += (6,8); X[6,8] = 1;
    SD += (7,8); X[7,8] = 1;
  }

  proc init(verbose:bool) {
    super.init(verbose=verbose);
    this.complete();
  }

  proc testMatrixOperators() {
    var vn2: [1..0] string;
    for n in vn do vn2.push_back(n);
    //vn2.push_back("nebula");
    var nm = new NamedMatrix(X=X, names=vn2);
    nm.set(i=2, j=7, w=13.1);
    assertRealEquals("Can set a new entry in matrix", expected=13.1, actual=nm.get(2,7));
    nm.remove(i=2, j=7);
    assertBoolEquals("Can remove an entry from the matrix", expected=false, actual=nm.SD.member(2,7));

    // Same thing with names instead
    nm.set("star lord", "yondu", w=13.1);
    assertRealEquals("Can set a new entry in matrix by name", expected=13.1, actual=nm.get("star lord","yondu"));
    nm.remove("star lord", "yondu");
    assertBoolEquals("Can remove an entry from the matrix by name", expected=false, actual=nm.SD.member(2,7));

  }

  proc tropicalTesting() {
    var nv: int = 8,
        D: domain(2) = {1..nv, 1..nv},
        SD = CSRDomain(D),
        X: [SD] real;

    SD += (1,2); X[1,2] = 1;
    SD += (1,3); X[1,3] = 1;
    SD += (1,4); X[1,4] = 1;
    SD += (2,2); X[2,2] = 1;
    SD += (2,4); X[2,4] = 1;
    SD += (3,4); X[3,4] = 1;
    SD += (4,5); X[4,5] = 1;
    SD += (5,6); X[5,6] = 1;
    SD += (6,7); X[6,7] = 1;
    SD += (6,8); X[6,8] = 1;
    SD += (7,8); X[7,8] = 1;

    var X2 = tropic(X,X);
    writeln(X2);
    writeln(X2.domain);
    //writeln(&& reduce (X2 == X));
    /*
    var X3 = tropic(X2,X);
    writeln(X3);
    writeln(X3.domain);
    var X4 = tropic(X3,X);
    writeln(X4);
    writeln(X4.domain);
    var X5 = tropic(X4,X);
    writeln(X5);
    writeln(X5.domain);
    var X6 = tropic(X5,X);
    writeln(X6);
    writeln(X6.domain);*/

    writeln(tropicLimit(X));
/*
    if R.domain == X.domain {
      var same = && reduce (R == X);
      if same {
        writeln("Arrays Equal");
        writeln(R);
      } else {
        writeln("Arrays Not Equal");
        writeln(tropic(R,X));
      }
    } else {
      writeln("Domains Not Equal");
      writeln(tropic(R,X));
    }*/
  }

  proc testIndexSort() {
    var p: bool = true;
    const Arr = [7, 10, 23, 1];
    const Idx = [2.2, 3.3, 1.1, 4.4];

    var a: [1..4] real = for x in indexSort(arr=Arr, idx=Idx) do x;
    var expected: [1..4] real = [23, 7, 10, 1];
    assertArrayEquals(msg="IndexSort", expected=expected, actual=a);

    var b: [1..4] real = for x in indexSort(arr=Arr, idx=Idx, reverse=true) do x;
    var e: [1..4] real = [1, 10, 7, 23];
    assertArrayEquals(msg="IndexSort reversed", expected=e, actual=b);
  }

  proc testNamedMatrix() {
    write("testNamedMatrix() ...");
    var nm = new NamedMatrix(X=X);
    nm.set(1,3, 17.0);

    assertIntEquals(msg="nrows() set correctly", expected=8, actual=nm.nrows());
    assertIntEquals(msg="ncols() set correctly", expected=8, actual=nm.ncols());
    assertIntEquals(msg="Number of non-zeroes", expected=11:int, actual=11:int );
    assertRealEquals(msg="X.sparsity", expected=0.171875:real, actual=nm.sparsity());
    assertRealEquals(msg="Max of row 1", expected=17.0:real, actual=nm.rowMax(1));
    var e: [1..8] real = [17.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, NAN];
    assertArrayEquals(msg="Max of all rows", expected=e, actual=nm.rowMax());

    var DY = {1..3, 1..4},
        SDY = CSRDomain(DY),
        Y:[SDY] real;
    SDY +=(1,3); Y[1,3] = 1;
    var nm2 = new NamedMatrix(X=Y);
    assertIntEquals("grid2seq(tuple) works on the matrix", expected=9, actual=nm2.grid2seq((3,1)));
    assertIntEquals("grid2seq(i,j) works on the matrix", expected=9, actual=nm2.grid2seq(3,1));
    var sg = nm2.seq2grid(7);
    assertIntEquals("seq2grid works on the matrix", expected=2, actual=sg[1]);
    assertIntEquals("seq2grid works on the matrix", expected=3, actual=sg[2]);
    writeln("...done");
  }

  proc testSetRowNames() {
    var nm = new NamedMatrix(X=X);
    var vn2: [1..0] string;
    for n in vn do vn2.push_back(n);
    vn2.push_back("stakar ogord");

    try {
      nm.setRowNames(vn2);
      assertThrowsError(msg="Set Row Names wrong size NOT ENFORCED", passed=false, new Error());
    } catch err: DimensionMatchError {
      assertThrowsError(msg="Set Row Names wrong size ENFORCED",  passed=true, err=err);
    } catch err: Error {
      assertThrowsError(msg="Set Row Names wrong size ENFORCED",  passed=true, err=err);
    }

    try {
      nm.setRowNames(vn2);
      assertIntEquals(msg="Set Row Names right size ENFORCED",  expected=vn.size, actual=nm.nrows());
    } catch {
      assertIntEquals(msg="Set Row Names right size failed",  expected=8, actual=nm.nrows());
    }
  }

  proc testSetColNames() {
    var nm = new NamedMatrix(X=X);
    try {
      nm.setColNames(vn);
      assertIntEquals(msg="Set Col Names right size ENFORCED",  expected=vn.size, actual=nm.ncols());
    } catch {
      assertIntEquals(msg="Set Col Names right size",  expected=vn.size, actual=nm.ncols());
    }
  }

  proc testNamedMatrixInitWithNames() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      assertIntEquals(msg="Initialize X with names only (rows)", expected=vn.size, actual=nm.nrows());
      assertIntEquals(msg="Initialize X with names only (cols)", expected=vn.size, actual=nm.ncols());
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc testSetByName() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      nm.set("star lord", "yondu", 17.0);
      assertRealEquals(msg="Set/Get by name", expected=17.0, actual=nm.get("star lord", "yondu"));
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc testArgMax() {
    var x: [1..3] real = [1.1, 3.3, 2.2];
    var y: [1..3,1..3] real = ((1,0,0), (0,0,2), (0,3,0));
    var xSD: sparse subdomain(x.domain);
    var z: [xSD] real;
    xSD += 2; z[2] = 1.3;
    xSD += 3; z[3] = 0.8;

    assertBoolEquals("z is sparse", expected=true, actual=isSparseArr(z));
    assertIntEquals("argmax1d(z)", expected=2, actual=argmax1d(z));

    assertIntEquals("argmax(x)", expected=2, actual=argmax(x));
    assertIntArrayEquals("argmax(y, ?)", expected=[3,2], actual=argmax(y));
    assertIntArrayEquals("argmax(y, ?)", expected=[3,2], actual=argmax(y));
    assertIntArrayEquals("argmax(y,0)", expected=[3,2], actual=argmax(y, axis=0));
    assertIntArrayEquals("argmax(y,1)", expected=[1,3,2], actual=argmax(y, axis=1));
    assertIntArrayEquals("argmax(y,2)", expected=[1,3,2], actual=argmax(y, axis=2));

    var SD2: sparse subdomain(D),
        X2:[SD2] real;
    SD2 += (1,2); X2[1,2] = 1;
    SD2 += (3,1); X2[3,1] = 3;
    SD2 += (3,4); X2[3,4] = 4;
    SD2 += (3,6); X2[3,6] = 1;
    var m = new NamedMatrix(X=X2, names=vn);
    assertRealEquals("m rowMax by row number", expected=4, actual=m.rowMax(3));
    assertRealEquals("m rowMax by row name", expected=4, actual=m.rowMax("groot"));
    var e: [1..4] real = [3.0, 1.0, 4.0, 1.0];
    assertArrayEquals("m colMax", expected=e, actual=m.colMax());
    var f: [1..4] int = [1,2,3,4];
    //writeln(m.rowArgMax());
    //assertIntArrayEquals("m rowArgMax", expected=f, actual=m.rowArgMax());
    //assertIntEquals("m rowArgMax(3)", expected=4, actual=m.rowArgMax(3));
    /* This interface is kind of annoying, I know what row I'm sending in so I
       should just get the answer back */
    assertIntEquals("m rowArgMax(3)", expected=4, actual=m.rowArgMax(3)[2]);
    assertIntEquals("m rowArgMax('groot')", expected=4, actual=m.rowArgMax('groot')[2]);

    assertIntEquals("m rowArgMin(3)", expected=6, actual=m.rowArgMin(3)[2]);
    assertIntEquals("m rowArgMin('groot')", expected=6, actual=m.rowArgMin('groot')[2]);

    assertRealEquals("m rowMin(3)", expected=1, actual=m.rowMin(3));
    assertRealEquals("m rowMin('groot')", expected=1, actual=m.rowMin('groot'));

    assertRealEquals("m colMin(6)", expected=1, actual=m.colMin(6));
    assertRealEquals("m colMin('mantis')", expected=1, actual=m.colMin('mantis'));

    assertRealEquals("m colMax(6)", expected=1, actual=m.colMax(6));
    assertRealEquals("m colMax('mantis')", expected=1, actual=m.colMax('mantis'));

    assertIntEquals("m colArgMax(6)", expected=6, actual=m.colArgMax(6)[2]);
    assertIntEquals("m colArgMax('mantis')", expected=6, actual=m.colArgMax('mantis')[2]);

    assertIntEquals("m colArgMin(6)", expected=6, actual=m.colArgMin(6)[2]);
    assertIntEquals("m colArgMin('mantis')", expected=6, actual=m.colArgMin('mantis')[2]);

  }

  proc testCosineDistance() {
    var X = Matrix(
      [3.0, 0.2, 0.0, 0.7, 0.1],
      [0.2, 2.0, 0.3, 0.0, 0.0],
      [0.0, 0.3, 3.0, 0.9, 0.6],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.1, 0.0, 0.6, 0.0, 2.0]
    );

    var Y = Matrix(
      [3.0, 0.2, 0.0, 0.7, 0.1],
      [0.2, 2.0, 0.3, 0.0, 0.0],
      [0.0, 0.3, 3.0, 0.9, 0.6],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.1, 0.0, 0.6, 0.0, 2.0]
    );

    const xxTarget = Matrix(
      [0.0, 0.974619, 0.992338, 0.930778, 0.988007],
      [0.974619, 0.0, 0.964601, 0.981269, 0.988919],
      [0.992338, 0.964601, 0.0, 0.917246, 0.93309],
      [0.930778, 0.981269, 0.917246, 0.0, 0.973663],
      [0.988007, 0.988919, 0.93309, 0.973663, 0.0]);

    var V = cosineDistance(X);
    const a = V-xxTarget;
    const aa = norm(a);
    assertRealApproximates(msg="Cosine Distance(X) norm", expected=1.39884e-06, actual=aa);

    var V2 = cosineDistance(X,Y);
    const cosimXYtarget = Matrix(
      [0.895178, 0.974619, 0.992338, 0.930778, 0.930778, 0.988007],
      [0.974619, 0.757869, 0.964601, 0.981269, 0.981269, 0.988919],
      [0.992338, 0.964601, 0.902534, 0.917246, 0.917246, 0.93309],
      [0.930778, 0.981269, 0.917246, 0.811321, 0.811321, 0.973663],
      [0.988007, 0.988919, 0.93309,  0.973663, 0.973663, 0.771167]
      );
    const b = V2-cosimXYtarget;
    const bb = norm(b);
    assertRealApproximates(msg="Cosine Distance(X,Y) norm", expected=1.53583e-06, actual=bb);
  }

  /*
   Does not assert anything, LabeledMatrix may be deprecated
   */
  proc testLabelMatrix() {
    var Y = Matrix(
      [3.0, 0.2, 0.0, 0.7, 0.1],
      [0.2, 2.0, 0.3, 0.0, 0.0],
      [0.0, 0.3, 3.0, 0.9, 0.6],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.1, 0.0, 0.6, 0.0, 2.0]
    );
    var L = new LabelMatrix();
    L.fromMatrix(Y);
    assertRealEquals(msg="Label Matrix entry (1,4)", expected=0.7, actual=L.data(1,4));
  }

  proc testECDF() {
    var x = [3,3,1,4];
    var ecdf = new ECDF(x);
    assertIntEquals("ECDF: Number of observations", expected=4, actual=ecdf.nobs);

    var y = [3.0, 55.0, 0.5, 1.5];
    var d = ecdf(y);
    assertArrayEquals("ECDF: Output", expected=[0.75,1.0,0.0,0.25], actual=d);
  }

  proc testBiMap() {
    var bm = new BiMap();
    bm.add("bob");
    bm.add("chuck");
    bm.add("bob");
    bm.add('ethel', 78);
    bm.add('frank', 3);

    assertIntEquals("BIMAP redundant entry ignored", expected=4, actual=bm.keys.size);
    for k in bm.keys {
      assertStringEquals("BIMAP name is retrieved", k, bm.idx[bm.ids[k]]);
      assertIntEquals("BIMAP id is retrieved", bm.ids[k], bm.get(k));
    }
    assertIntEquals("BIMPA max is set", expected=78, actual=bm.max());

    var abm = new BiMap();
    abm.add("one",1);
    abm.add("two",2);
    var bbm = new BiMap();
    bbm.add("three",3);
    bbm.add("four",4);

    var cbm = abm.uni(bbm);
    assertIntEquals("BIMAP union size", expected=4, actual=cbm.keys.size);
    assertIntEquals("BIMAP union max", expected=4, actual=cbm.max());
  }

  proc testViterbi() {
    var obs = ["normal", "cold", "dizzy"];
    var states = ["health", "fever"];
    var initP = [0.6, 0.4];

    /* (Hidden) State transition matrix
             Healthy Fever
    Healthy  0.7     0.3
      Fever  0.4     0.6
     */
    var A: [1..2, 1..2] real = Matrix(
      [0.7, 0.3],
      [0.4, 0.6]
    );

    /*
     State emission probabilities
             Normal Cold Dizzy
    Healthy  0.5    0.4  0.1
      Fever  0.1    0.3  0.6
     */
     var B: [1..2, 1..3] real = Matrix(
       [0.5, 0.4, 0.1],
       [0.1, 0.3, 0.6]
    );
    var v = Viterbi(obs, states, initP, transitionProbabilities=A, emissionProbabilities=B);
    assertStringArrayEquals("Viterbi states output", expected=["health","health","fever"], actual=v);
  }

  proc testLetters() {
    var ltrs = letters(28);
    assertStringEquals("First in letters is A", expected="A", actual=ltrs[1]);
    assertStringEquals("28th in letters is BB", expected="BB", actual=ltrs[28]);

    var gn = gridNames(7);
    assertIntEquals("Grid names(7) has 49 entries", expected=49, actual=gn.size);
    assertStringEquals("Last grid name is 'G7'", expected="G7", actual=gn[49]);
    var gn2 = gridNames(i=3, j=5);
    assertRealEquals("gridNames(3,5) has 15 entries", expected=15, actual=gn2.size);
    assertStringEquals("Last grid name is 'C5'", expected="C5", actual=gn2[15]);
  }

  proc testChoice() {
    var x = [1,2,3,4,5,6,7,8,9,10];
    var y = choice(a=x,size=2);
    assertIntEquals("Can pick two choices, no replacement, no p", expected=2, actual=y.size);
    var z = choice(x,size=2,replace=true);
    assertIntEquals("Can pick two choices, with replacement, no p", expected=2, actual=z.size);

    var s = [1, 2, 3, 4, 5];
    // Note, this does not need to be normalized but it is easier to see
    var p = [0.5, 0.25, 0.13, 0.09, 0.03];
    var m = chooseMultinomial(a=s, replace=false, size=3, p=p);
    assertIntEquals("Multinomial w/o replace returns correct number of results", expected=3, actual=m.size);
    var n = chooseMultinomial(a=s, replace=true, size=3, p=p);
    assertIntEquals("Multinomial w/ replace returns correct number of results", expected=3, actual=n.size);
    var o = choice(a=s, replace=false, size=3, p=p);
  }

  proc run() {
    super.run();
    testMatrixOperators();
    //tropicalTesting();
    testIndexSort();
    testNamedMatrix();
    testSetRowNames();
    testSetColNames();
    testNamedMatrixInitWithNames();
    testSetByName();
    testArgMax();
    testCosineDistance();
    testLabelMatrix();
    testECDF();
    testBiMap();
    testViterbi();
    testLetters();
    testChoice();
    return 0;
  }
}

proc main(args: [] string) : int {
  var t = new NumSuchTest(verbose=false);
  var ret = t.run();
  t.report();
  return ret;
}

use NumSuch,
    Norm,
    LinearAlgebra,
    Viterbi,
    Core,
    Charcoal;

class NumSuchTest : UnitTest {
  var nv: int = 8,
      D: domain(2) = {1..nv, 1..nv},
      SD: sparse subdomain(D),
      X: [SD] real;

  var vn: [1..0] string;

  proc setUp() {
    vn.push_back("star lord");
    vn.push_back("gamora");
    vn.push_back("groot");
    vn.push_back("drax");
    vn.push_back("rocket");
    vn.push_back("mantis");
    vn.push_back("yondu");
    vn.push_back("nebula");

    SD += (1,2); X[1,2] = 1;
    SD += (1,3); X[1,3] = 1;
    SD += (1,4); X[1,4] = 1;
    SD += (2,2); X[2,2] = 1;
    SD += (2,4); X[2,4] = 1;
    SD += (3,4); X[3,4] = 1;
    SD += (4,5); X[4,5] = 1;
    SD += (5,6); X[5,6] = 1;
    SD += (6,7); X[6,7] = 1;
    SD += (6,8); X[6,8] = 1;
    SD += (7,8); X[7,8] = 1;
  }

  proc init(verbose:bool) {
    super.init(verbose=verbose);
    this.complete();
  }

  proc testMatrixOperators() {
    var vn2: [1..0] string;
    for n in vn do vn2.push_back(n);
    //vn2.push_back("nebula");
    var nm = new NamedMatrix(X=X, names=vn2);
    nm.set(i=2, j=7, w=13.1);
    assertRealEquals("Can set a new entry in matrix", expected=13.1, actual=nm.get(2,7));
    nm.remove(i=2, j=7);
    assertBoolEquals("Can remove an entry from the matrix", expected=false, actual=nm.SD.member(2,7));

    // Same thing with names instead
    nm.set("star lord", "yondu", w=13.1);
    assertRealEquals("Can set a new entry in matrix by name", expected=13.1, actual=nm.get("star lord","yondu"));
    nm.remove("star lord", "yondu");
    assertBoolEquals("Can remove an entry from the matrix by name", expected=false, actual=nm.SD.member(2,7));

  }

  proc tropicalTesting() {
    var nv: int = 8,
        D: domain(2) = {1..nv, 1..nv},
        SD = CSRDomain(D),
        X: [SD] real;

    SD += (1,2); X[1,2] = 1;
    SD += (1,3); X[1,3] = 1;
    SD += (1,4); X[1,4] = 1;
    SD += (2,2); X[2,2] = 1;
    SD += (2,4); X[2,4] = 1;
    SD += (3,4); X[3,4] = 1;
    SD += (4,5); X[4,5] = 1;
    SD += (5,6); X[5,6] = 1;
    SD += (6,7); X[6,7] = 1;
    SD += (6,8); X[6,8] = 1;
    SD += (7,8); X[7,8] = 1;

    var X2 = tropic(X,X);
    writeln(X2);
    writeln(X2.domain);
    //writeln(&& reduce (X2 == X));
    /*
    var X3 = tropic(X2,X);
    writeln(X3);
    writeln(X3.domain);
    var X4 = tropic(X3,X);
    writeln(X4);
    writeln(X4.domain);
    var X5 = tropic(X4,X);
    writeln(X5);
    writeln(X5.domain);
    var X6 = tropic(X5,X);
    writeln(X6);
    writeln(X6.domain);*/

    writeln(tropicLimit(X));
/*
    if R.domain == X.domain {
      var same = && reduce (R == X);
      if same {
        writeln("Arrays Equal");
        writeln(R);
      } else {
        writeln("Arrays Not Equal");
        writeln(tropic(R,X));
      }
    } else {
      writeln("Domains Not Equal");
      writeln(tropic(R,X));
    }*/
  }

  proc testIndexSort() {
    var p: bool = true;
    const Arr = [7, 10, 23, 1];
    const Idx = [2.2, 3.3, 1.1, 4.4];

    var a: [1..4] real = for x in indexSort(arr=Arr, idx=Idx) do x;
    var expected: [1..4] real = [23, 7, 10, 1];
    assertArrayEquals(msg="IndexSort", expected=expected, actual=a);

    var b: [1..4] real = for x in indexSort(arr=Arr, idx=Idx, reverse=true) do x;
    var e: [1..4] real = [1, 10, 7, 23];
    assertArrayEquals(msg="IndexSort reversed", expected=e, actual=b);
  }

  proc testNamedMatrix() {
    write("testNamedMatrix() ...");
    var nm = new NamedMatrix(X=X);
    nm.set(1,3, 17.0);

    assertIntEquals(msg="nrows() set correctly", expected=8, actual=nm.nrows());
    assertIntEquals(msg="ncols() set correctly", expected=8, actual=nm.ncols());
    assertIntEquals(msg="Number of non-zeroes", expected=11:int, actual=11:int );
    assertRealEquals(msg="X.sparsity", expected=0.171875:real, actual=nm.sparsity());
    assertRealEquals(msg="Max of row 1", expected=17.0:real, actual=nm.rowMax(1));
    var e: [1..8] real = [17.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, NAN];
    assertArrayEquals(msg="Max of all rows", expected=e, actual=nm.rowMax());
    assertIntEquals("grid2seq(tuple) works on the matrix", expected=23, actual=nm.grid2seq((3,7)));
    assertIntEquals("grid2seq(i,j) works on the matrix", expected=23, actual=nm.grid2seq(3,7));
    var sg = nm.seq2grid(23);
    assertIntEquals("seq2grid works on the matrix", expected=3, actual=sg[1]);
    assertIntEquals("seq2grid works on the matrix", expected=7, actual=sg[2]);
    writeln("...done");
  }

  proc testSetRowNames() {
    var nm = new NamedMatrix(X=X);
    var vn2: [1..0] string;
    for n in vn do vn2.push_back(n);
    vn2.push_back("stakar ogord");

    try {
      nm.setRowNames(vn2);
      assertThrowsError(msg="Set Row Names wrong size NOT ENFORCED", passed=false, new Error());
    } catch err: DimensionMatchError {
      assertThrowsError(msg="Set Row Names wrong size ENFORCED",  passed=true, err=err);
    } catch err: Error {
      assertThrowsError(msg="Set Row Names wrong size ENFORCED",  passed=true, err=err);
    }

    try {
      nm.setRowNames(vn2);
      assertIntEquals(msg="Set Row Names right size ENFORCED",  expected=vn.size, actual=nm.nrows());
    } catch {
      assertIntEquals(msg="Set Row Names right size failed",  expected=8, actual=nm.nrows());
    }
  }

  proc testSetColNames() {
    var nm = new NamedMatrix(X=X);
    try {
      nm.setColNames(vn);
      assertIntEquals(msg="Set Col Names right size ENFORCED",  expected=vn.size, actual=nm.ncols());
    } catch {
      assertIntEquals(msg="Set Col Names right size",  expected=vn.size, actual=nm.ncols());
    }
  }

  proc testNamedMatrixInitWithNames() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      assertIntEquals(msg="Initialize X with names only (rows)", expected=vn.size, actual=nm.nrows());
      assertIntEquals(msg="Initialize X with names only (cols)", expected=vn.size, actual=nm.ncols());
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc testSetByName() {
    try {
      var nm = new NamedMatrix(rownames = vn, colnames=vn);
      nm.set("star lord", "yondu", 17.0);
      assertRealEquals(msg="Set/Get by name", expected=17.0, actual=nm.get("star lord", "yondu"));
    } catch e: DimensionMatchError {
      writeln(e);
    } catch e: Error {
      writeln(e);
    }
  }

  proc testArgMax() {
    var x: [1..3] real = [1.1, 3.3, 2.2];
    var y: [1..3,1..3] real = ((1,0,0), (0,0,2), (0,3,0));
    var xSD: sparse subdomain(x.domain);
    var z: [xSD] real;
    xSD += 2; z[2] = 1.3;
    xSD += 3; z[3] = 0.8;

    assertBoolEquals("z is sparse", expected=true, actual=isSparseArr(z));
    assertIntEquals("argmax1d(z)", expected=2, actual=argmax1d(z));

    assertIntEquals("argmax(x)", expected=2, actual=argmax(x));
    assertIntArrayEquals("argmax(y, ?)", expected=[3,2], actual=argmax(y));
    assertIntArrayEquals("argmax(y, ?)", expected=[3,2], actual=argmax(y));
    assertIntArrayEquals("argmax(y,0)", expected=[3,2], actual=argmax(y, axis=0));
    assertIntArrayEquals("argmax(y,1)", expected=[1,3,2], actual=argmax(y, axis=1));
    assertIntArrayEquals("argmax(y,2)", expected=[1,3,2], actual=argmax(y, axis=2));

    var SD2: sparse subdomain(D),
        X2:[SD2] real;
    SD2 += (1,2); X2[1,2] = 1;
    SD2 += (3,1); X2[3,1] = 3;
    SD2 += (3,4); X2[3,4] = 4;
    SD2 += (3,6); X2[3,6] = 1;
    var m = new NamedMatrix(X=X2, names=vn);
    assertRealEquals("m rowMax by row number", expected=4, actual=m.rowMax(3));
    assertRealEquals("m rowMax by row name", expected=4, actual=m.rowMax("groot"));
    var e: [1..4] real = [3.0, 1.0, 4.0, 1.0];
    assertArrayEquals("m colMax", expected=e, actual=m.colMax());
    var f: [1..4] int = [1,2,3,4];
    //writeln(m.rowArgMax());
    //assertIntArrayEquals("m rowArgMax", expected=f, actual=m.rowArgMax());
    //assertIntEquals("m rowArgMax(3)", expected=4, actual=m.rowArgMax(3));
    /* This interface is kind of annoying, I know what row I'm sending in so I
       should just get the answer back */
    assertIntEquals("m rowArgMax(3)", expected=4, actual=m.rowArgMax(3)[2]);
    assertIntEquals("m rowArgMax('groot')", expected=4, actual=m.rowArgMax('groot')[2]);

    assertIntEquals("m rowArgMin(3)", expected=6, actual=m.rowArgMin(3)[2]);
    assertIntEquals("m rowArgMin('groot')", expected=6, actual=m.rowArgMin('groot')[2]);

    assertRealEquals("m rowMin(3)", expected=1, actual=m.rowMin(3));
    assertRealEquals("m rowMin('groot')", expected=1, actual=m.rowMin('groot'));

    assertRealEquals("m colMin(6)", expected=1, actual=m.colMin(6));
    assertRealEquals("m colMin('mantis')", expected=1, actual=m.colMin('mantis'));

    assertRealEquals("m colMax(6)", expected=1, actual=m.colMax(6));
    assertRealEquals("m colMax('mantis')", expected=1, actual=m.colMax('mantis'));

    assertIntEquals("m colArgMax(6)", expected=6, actual=m.colArgMax(6)[2]);
    assertIntEquals("m colArgMax('mantis')", expected=6, actual=m.colArgMax('mantis')[2]);

    assertIntEquals("m colArgMin(6)", expected=6, actual=m.colArgMin(6)[2]);
    assertIntEquals("m colArgMin('mantis')", expected=6, actual=m.colArgMin('mantis')[2]);

  }

  proc testCosineDistance() {
    var X = Matrix(
      [3.0, 0.2, 0.0, 0.7, 0.1],
      [0.2, 2.0, 0.3, 0.0, 0.0],
      [0.0, 0.3, 3.0, 0.9, 0.6],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.1, 0.0, 0.6, 0.0, 2.0]
    );

    var Y = Matrix(
      [3.0, 0.2, 0.0, 0.7, 0.1],
      [0.2, 2.0, 0.3, 0.0, 0.0],
      [0.0, 0.3, 3.0, 0.9, 0.6],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.1, 0.0, 0.6, 0.0, 2.0]
    );

    const xxTarget = Matrix(
      [0.0, 0.974619, 0.992338, 0.930778, 0.988007],
      [0.974619, 0.0, 0.964601, 0.981269, 0.988919],
      [0.992338, 0.964601, 0.0, 0.917246, 0.93309],
      [0.930778, 0.981269, 0.917246, 0.0, 0.973663],
      [0.988007, 0.988919, 0.93309, 0.973663, 0.0]);

    var V = cosineDistance(X);
    const a = V-xxTarget;
    const aa = norm(a);
    assertRealApproximates(msg="Cosine Distance(X) norm", expected=1.39884e-06, actual=aa);

    var V2 = cosineDistance(X,Y);
    const cosimXYtarget = Matrix(
      [0.895178, 0.974619, 0.992338, 0.930778, 0.930778, 0.988007],
      [0.974619, 0.757869, 0.964601, 0.981269, 0.981269, 0.988919],
      [0.992338, 0.964601, 0.902534, 0.917246, 0.917246, 0.93309],
      [0.930778, 0.981269, 0.917246, 0.811321, 0.811321, 0.973663],
      [0.988007, 0.988919, 0.93309,  0.973663, 0.973663, 0.771167]
      );
    const b = V2-cosimXYtarget;
    const bb = norm(b);
    assertRealApproximates(msg="Cosine Distance(X,Y) norm", expected=1.53583e-06, actual=bb);
  }

  /*
   Does not assert anything, LabeledMatrix may be deprecated
   */
  proc testLabelMatrix() {
    var Y = Matrix(
      [3.0, 0.2, 0.0, 0.7, 0.1],
      [0.2, 2.0, 0.3, 0.0, 0.0],
      [0.0, 0.3, 3.0, 0.9, 0.6],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.7, 0.0, 0.9, 2.0, 0.0],
      [0.1, 0.0, 0.6, 0.0, 2.0]
    );
    var L = new LabelMatrix();
    L.fromMatrix(Y);
    assertRealEquals(msg="Label Matrix entry (1,4)", expected=0.7, actual=L.data(1,4));
  }

  proc testECDF() {
    var x = [3,3,1,4];
    var ecdf = new ECDF(x);
    assertIntEquals("ECDF: Number of observations", expected=4, actual=ecdf.nobs);

    var y = [3.0, 55.0, 0.5, 1.5];
    var d = ecdf(y);
    assertArrayEquals("ECDF: Output", expected=[0.75,1.0,0.0,0.25], actual=d);
  }

  proc testBiMap() {
    var bm = new BiMap();
    bm.add("bob");
    bm.add("chuck");
    bm.add("bob");
    bm.add('ethel', 78);
    bm.add('frank', 3);

    assertIntEquals("BIMAP redundant entry ignored", expected=4, actual=bm.keys.size);
    for k in bm.keys {
      assertStringEquals("BIMAP name is retrieved", k, bm.idx[bm.ids[k]]);
      assertIntEquals("BIMAP id is retrieved", bm.ids[k], bm.get(k));
    }
    assertIntEquals("BIMPA max is set", expected=78, actual=bm.max());

    var abm = new BiMap();
    abm.add("one",1);
    abm.add("two",2);
    var bbm = new BiMap();
    bbm.add("three",3);
    bbm.add("four",4);

    var cbm = abm.uni(bbm);
    assertIntEquals("BIMAP union size", expected=4, actual=cbm.keys.size);
    assertIntEquals("BIMAP union max", expected=4, actual=cbm.max());
  }

  proc testViterbi() {
    var obs = ["normal", "cold", "dizzy"];
    var states = ["health", "fever"];
    var initP = [0.6, 0.4];

    /* (Hidden) State transition matrix
             Healthy Fever
    Healthy  0.7     0.3
      Fever  0.4     0.6
     */
    var A: [1..2, 1..2] real = Matrix(
      [0.7, 0.3],
      [0.4, 0.6]
    );

    /*
     State emission probabilities
             Normal Cold Dizzy
    Healthy  0.5    0.4  0.1
      Fever  0.1    0.3  0.6
     */
     var B: [1..2, 1..3] real = Matrix(
       [0.5, 0.4, 0.1],
       [0.1, 0.3, 0.6]
    );
    var v = Viterbi(obs, states, initP, transitionProbabilities=A, emissionProbabilities=B);
    assertStringArrayEquals("Viterbi states output", expected=["health","health","fever"], actual=v);
  }

  proc testLetters() {
    var ltrs = letters(28);
    assertStringEquals("First in letters is A", expected="A", actual=ltrs[1]);
    assertStringEquals("28th in letters is BB", expected="BB", actual=ltrs[28]);

    var gn = gridNames(7);
    assertIntEquals("Grid names(7) has 49 entries", expected=49, actual=gn.size);
    assertStringEquals("Last grid name is 'G7'", expected="G7", actual=gn[49]);
    var gn2 = gridNames(i=3, j=5);
    assertRealEquals("gridNames(3,5) has 15 entries", expected=15, actual=gn2.size);
    assertStringEquals("Last grid name is 'C5'", expected="C5", actual=gn2[15]);
  }

  proc testChoice() {
    var x = [1,2,3,4,5,6,7,8,9,10];
    var y = choice(a=x,size=2);
    assertIntEquals("Can pick two choices, no replacement, no p", expected=2, actual=y.size);
    var z = choice(x,size=2,replace=true);
    assertIntEquals("Can pick two choices, with replacement, no p", expected=2, actual=z.size);

    var s = [1, 2, 3, 4, 5];
    // Note, this does not need to be normalized but it is easier to see
    var p = [0.5, 0.25, 0.13, 0.09, 0.03];
    var m = chooseMultinomial(a=s, replace=false, size=3, p=p);
    assertIntEquals("Multinomial w/o replace returns correct number of results", expected=3, actual=m.size);
    var n = chooseMultinomial(a=s, replace=true, size=3, p=p);
    assertIntEquals("Multinomial w/ replace returns correct number of results", expected=3, actual=n.size);
    var o = choice(a=s, replace=false, size=3, p=p);
  }

  proc run() {
    super.run();
    testMatrixOperators();
    //tropicalTesting();
    testIndexSort();
    testNamedMatrix();
    testSetRowNames();
    testSetColNames();
    testNamedMatrixInitWithNames();
    testSetByName();
    testArgMax();
    testCosineDistance();
    testLabelMatrix();
    testECDF();
    testBiMap();
    testViterbi();
    testLetters();
    testChoice();
    return 0;
  }
}

proc main(args: [] string) : int {
  var t = new NumSuchTest(verbose=false);
  var ret = t.run();
  t.report();
  return ret;
}


Add some returns [approved by Michael in chat]

Through my default initializer work, I discovered that the compiler had failed
to complain about buffer.copyout and buffer.copyin specifying their return type
but not actually returning anything.  These methods should return the result of
their calls to the other overload of themselves, so I'm adding that return (and
investigating why the compiler didn't complain before).

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@23762 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

module BufferInternals {
  use SysBasic;
  use Error;

  extern type qbytes_ptr_t;
  extern type qbuffer_ptr_t;
  extern type qbuffer_iter_t;
  extern const QBYTES_PTR_NULL:qbytes_ptr_t;
  extern const QBUFFER_PTR_NULL:qbuffer_ptr_t;

  extern proc qbytes_retain(qb:qbytes_ptr_t);
  extern proc qbytes_release(qb:qbytes_ptr_t);
  extern proc qbytes_len(qb:qbytes_ptr_t):int(64);
  extern proc qbytes_data(qb:qbytes_ptr_t):c_void_ptr;

  extern proc qbytes_create_iobuf(out ret:qbytes_ptr_t):err_t;
  extern proc qbytes_create_calloc(out ret:qbytes_ptr_t, len:int(64)):err_t;

  extern proc qbuffer_iter_null():qbuffer_iter_t;

  /*
  extern proc qbuffer_init(ref buf:qbuffer_t):err_t;
  extern proc qbuffer_destroy(ref buf:qbuffer_t);
  */
  extern proc qbuffer_create(out buf:qbuffer_ptr_t):err_t;
  extern proc qbuffer_retain(buf:qbuffer_ptr_t);
  extern proc qbuffer_release(buf:qbuffer_ptr_t);

  extern proc qbuffer_append(buf:qbuffer_ptr_t, bytes:qbytes_ptr_t, skip_bytes:int(64), len_bytes:int(64)):err_t;
  extern proc qbuffer_append_buffer(buf:qbuffer_ptr_t, src:qbuffer_ptr_t, src_start:qbuffer_iter_t, src_end:qbuffer_iter_t):err_t;
  extern proc qbuffer_prepend(buf:qbuffer_ptr_t, bytes:qbytes_ptr_t, skip_bytes:int(64), len_bytes:int(64)):err_t;
  extern proc qbuffer_flatten(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, out bytes_out):err_t;
  //extern proc qbuffer_clone(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, out buf_out:qbuffer_t):err_t;
  extern proc qbuffer_copyout(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, ref x, size):err_t;
  extern proc qbuffer_copyin(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, ref x, size):err_t;

  extern proc qbuffer_begin(buf:qbuffer_ptr_t):qbuffer_iter_t;
  extern proc qbuffer_end(buf:qbuffer_ptr_t):qbuffer_iter_t;
  extern proc qbuffer_iter_next_part(buf:qbuffer_ptr_t, ref it:qbuffer_iter_t);
  extern proc qbuffer_iter_prev_part(buf:qbuffer_ptr_t, ref it:qbuffer_iter_t);
  extern proc qbuffer_iter_advance(buf:qbuffer_ptr_t, ref it:qbuffer_iter_t, amt:int(64));

  extern proc qbuffer_iter_get(it: qbuffer_iter_t, end:qbuffer_iter_t,
                                out bytes_out:qbytes_ptr_t,
                                out skip_out:int(64),
                                out len_out:int(64));
  extern proc qbuffer_iter_num_bytes(start:qbuffer_iter_t, end:qbuffer_iter_t):int(64);

  extern proc qbuffer_len(buf:qbuffer_ptr_t):int(64);

  extern proc debug_print_qbuffer_iter(/*const*/ ref it:qbuffer_iter_t);

  extern proc qbuffer_start_offset(buf:qbuffer_ptr_t):int(64);
  extern proc qbuffer_end_offset(buf:qbuffer_ptr_t):int(64);
  extern proc qbuffer_reposition(buf:qbuffer_ptr_t, new_offset_start:int(64));

  extern proc bulk_get_bytes(src_locale:int, src_addr:qbytes_ptr_t):qbytes_ptr_t;

  extern proc bulk_put_buffer(dst_locale:int, dst_addr:c_void_ptr, dst_len:int(64), buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t):err_t;

}

module Buffers {
  use BufferInternals; // TODO -- non-exporting use

  // Now define the Chapel types using the originals..
  record bytes {
    var home: locale;
    var _bytes_internal:qbytes_ptr_t;
  }

  record buffer_iterator {
    var home: locale;
    var _bufit_internal:qbuffer_iter_t;
  }
  record buffer_range {
    var start:buffer_iterator;
    var end:buffer_iterator;
  }

  proc buffer_range.len:int(64) {
    var ret:int(64);
    on this.start.home {
      ret = qbuffer_iter_num_bytes(this.start._bufit_internal,
                                   this.end._bufit_internal);
    }
    return ret;
  }

  // FUTURE -- we could create const buffers
  // with a boolean type param
  record buffer {
    var home:locale;
    var _buf_internal:qbuffer_ptr_t;
  }

  // bytes methods.
  proc bytes.bytes() {
    this._bytes_internal = QBYTES_PTR_NULL;
    this.home = here;
  }
  proc bytes.bytes(len:int(64), out error:syserr) {
    error = qbytes_create_calloc(this._bytes_internal, len);
    // The buffer is "retained" internally on creation, but only on success.
    this.home = here;
  }
  proc bytes.bytes(len:int(64)) {
    var error:syserr = ENOERR;
    error = qbytes_create_calloc(this._bytes_internal, len);
    if error then ioerror(error, "in bytes constructor");
    // The buffer is retained internally on construction, but only on success.
    this.home = here;
  }

  proc create_iobuf(out error:syserr):bytes {
    var ret:bytes;
    error = qbytes_create_iobuf(ret._bytes_internal);
    // The buffer is "retained" internally on creation, but only on success.
    ret.home = here;
    return ret;
  }
  proc create_iobuf():bytes {
    var err:syserr = ENOERR;
    var ret = create_iobuf(err);
    if err then ioerror(err, "in create_iobuf");
    return ret;
  }

  // TODO -- shouldn't have to write this this way!
  pragma "init copy fn"
  proc chpl__initCopy(x: bytes) {
    if x.home == here {
      qbytes_retain(x._bytes_internal);
      return x;
    } else {
      var ret:bytes;
      ret.home = here;
      writeln("Bulk moving bytes");
      // The initial ref count is 1, so no need to call qbytes_retain here.
      ret._bytes_internal = bulk_get_bytes(x.home.id, x._bytes_internal);
      return ret;
    }
  }

  proc =(ref ret:bytes, x:bytes) {
    // retain -- release
    if( x.home == here ) {
      on x.home {
        qbytes_retain(x._bytes_internal);
      }

      on ret.home {
        qbytes_release(ret._bytes_internal);
      }

      ret.home = x.home;
      ret._bytes_internal = x._bytes_internal;
    } else {
      on ret.home {
        qbytes_release(ret._bytes_internal);
      }
      ret.home = here;
      writeln("Bulk moving bytes");
      ret._bytes_internal = bulk_get_bytes(x.home.id, x._bytes_internal);
      // On return from bulk_get_bytes, the ref count in ret._bytes_internal
      // should be 1.
      // Note that the error case is not handled (bulk_get_bytes must succeed).
    }
  }

  proc bytes.~bytes() {
    on this.home {
      qbytes_release(this._bytes_internal);
      this._bytes_internal = QBYTES_PTR_NULL;
    }
  }
  proc bytes.len:int(64) {
    var ret:int(64);
    on home {
     ret = qbytes_len(this._bytes_internal);
    }
    return ret;
  }

  // buffer iterator
  proc buffer_iterator.buffer_iterator() {
    this.home = here;
    this._bufit_internal = qbuffer_iter_null();
  }
  /* don't believe this is necessary...
  pragma "init copy fn"
  proc chpl__initCopy(x: buffer_iterator) {
    return x;
  }
  */

  // buffer methods.
  proc buffer.buffer(out error:syserr) {
    this.home = here;
    error = qbuffer_create(this._buf_internal);
  }
  proc buffer.buffer() {
    var error:syserr = ENOERR;
    this.home = here;
    error = qbuffer_create(this._buf_internal);
    if error then ioerror(error, "in buffer constructor");
  }

  // buffer.flatten is special in that it creates a
  // bytes object wherever it is called (and communicates it
  // intelligently). It is use in buffer's initCopy.
  proc buffer.flatten(range:buffer_range, out error:syserr):bytes {
    var ret:bytes;
    if this.home == here {
      error = qbuffer_flatten(this._buf_internal, range.start._bufit_internal, range.end._bufit_internal, ret);
    } else {
      var dst_locale = here;
      var dst_len:int(64) = range.len;
      ret = new bytes(dst_len);
      var dst_addr = qbytes_data(ret._bytes_internal);
      on this.home {
        // Copy the buffer to the bytes...
        error = bulk_put_buffer(dst_locale.id, dst_addr, dst_len,
                                this._buf_internal,
                                range.start._bufit_internal,
                                range.end._bufit_internal);
      }
    }
    return ret;
  }

  // TODO -- shouldn't have to write this this way!
  pragma "init copy fn"
  proc chpl__initCopy(x: buffer) {
    if x.home == here {
      qbuffer_retain(x._buf_internal);
      return x;
    } else {
      var ret:buffer;
      var start_offset:int(64);
      var end_offset:int(64);

      on x.home {
        start_offset = qbuffer_start_offset(x._buf_internal);
        end_offset = qbuffer_end_offset(x._buf_internal);
      }

      var b = new bytes(end_offset - start_offset);
      var len = qbytes_len(b._bytes_internal);
      var ptr = qbytes_data(b._bytes_internal);

      // Now create a local buffer with the right data
      // in it starting at the right position.
      qbuffer_reposition(ret._buf_internal, start_offset);

      var there_uid = here.id;

      on x.home {
        var err:syserr = ENOERR;
        err = bulk_put_buffer(there_uid, ptr, len, x._buf_internal,
                              qbuffer_begin(x._buf_internal),
                              qbuffer_end(x._buf_internal));
        if err then ioerror(err, "in buffer init copy");
      }

      ret.append(b);
      return ret;
    }
  }

  proc =(ref ret:buffer, x:buffer) {
    ret.home = here;
    // retain -- release
    if( x.home == ret.home ) {
      on x.home {
        qbuffer_retain(x._buf_internal);
      }

      on ret.home {
        qbuffer_release(ret._buf_internal);
      }

      ret.home = x.home;
      ret._buf_internal = x._buf_internal;
    } else {
      on ret.home {
        qbuffer_release(ret._buf_internal);
      }

      var start_offset:int(64);
      var end_offset:int(64);

      on x.home {
        start_offset = qbuffer_start_offset(x._buf_internal);
        end_offset = qbuffer_end_offset(x._buf_internal);
      }

      var b = new bytes(end_offset - start_offset);
      var len = qbytes_len(b._bytes_internal);
      var ptr = qbytes_data(b._bytes_internal);

      // Now create a local buffer with the right data
      // in it starting at the right position.
      qbuffer_reposition(ret._buf_internal, start_offset);

      var there_uid = here.id;

      on x.home {
        var err:syserr = ENOERR;
        err = bulk_put_buffer(there_uid, ptr, len, x._buf_internal,
                              qbuffer_begin(x._buf_internal),
                              qbuffer_end(x._buf_internal));
        if err then ioerror(err, "in buffer assignment");
      }

      ret.append(b);
    }
  }

  proc buffer.~buffer() {
    on this.home {
      qbuffer_release(this._buf_internal);
      this._buf_internal = QBUFFER_PTR_NULL;
    }
  }

  proc buffer.len {
    var ret:int(64);
    on this.home {
      ret = qbuffer_len(this._buf_internal);
    }
    return ret;
  }

  proc buffer.append(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len, out error:syserr) {
    on this.home {
      error = qbuffer_append(this._buf_internal, b._bytes_internal, skip_bytes, len_bytes);
    }
  }
  proc buffer.append(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len) {
    var err:syserr = ENOERR;
    this.append(b, skip_bytes, len_bytes, err);
    if err then ioerror(err, "in buffer.append");
  }

  proc buffer.append(buf:buffer, part:buffer_range = buf.all(), out error:syserr) {
    on this.home {
      error = qbuffer_append_buffer(this._buf_internal, buf._buf_internal, part.start._bufit_internal, part.end._bufit_internal);
    }
  }
  proc buffer.append(buf:buffer, part:buffer_range = buf.all()) {
    var err:syserr = ENOERR;
    this.append(buf, part, err);
    if err then ioerror(err, "in buffer.append");
  }

  proc buffer.prepend(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len, out error:syserr) {
    on this.home {
      qbuffer_prepend(this._buf_internal, b._bytes_internal, skip_bytes, len_bytes);
    }
  }
  proc buffer.prepend(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len) {
    var err:syserr = ENOERR;
    this.prepend(b, skip_bytes, len_bytes, err);
    if err then ioerror(err, "in buffer.prepend");
  }

  proc buffer.start():buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      ret._bufit_internal = qbuffer_begin(this._buf_internal);
    }
    return ret;
  }
  proc buffer.end():buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      ret._bufit_internal = qbuffer_end(this._buf_internal);
    }
    return ret;
  }
  proc buffer.all():buffer_range {
    return new buffer_range(this.start(), this.end());
  }

  proc buffer.next_part(inout it:buffer_iterator) {
    on this.home {
      var x = it;
      qbuffer_iter_next_part(this._buf_internal, x._bufit_internal);
      it = x;
    }
  }
  proc buffer.prev_part(inout it:buffer_iterator) {
    on this.home {
      var x = it;
      qbuffer_iter_prev_part(this._buf_internal, x._bufit_internal);
      it = x;
    }
  }
  proc buffer.advance(inout it:buffer_iterator, amount:int(64)) {
    on this.home {
      var x = it;
      qbuffer_iter_advance(this._buf_internal, x._bufit_internal, amount);
      it = x;
    }
  }

  proc buffer_iterator.debug_print()
  {
    on this.home {
      debug_print_qbuffer_iter(this._bufit_internal);
    }
  }

  /*
  proc buffer.clone(range:buffer_range, err:ErrorHandler = nil):buffer {
    on this.home {
      var ret:buffer;
      seterr(qbuffer_clone(this._buf_internal, range.start, range.end, ret._buf_internal));
      return ret;
    }
  }*/

  // methods to read/write basic types.
  proc buffer.copyout(it:buffer_iterator, out value, out error:syserr):buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      var end:buffer_iterator = it;
      var tmp:value.type;
      var sz = numBytes(value.type);
      this.advance(end, sz);
      error = qbuffer_copyout(this._buf_internal, it._bufit_internal, end._bufit_internal, tmp, sz);
      value = tmp;
      ret = end;
    }
    return ret;
  }
  proc buffer.copyout(it:buffer_iterator, out value):buffer_iterator {
    var err:syserr = ENOERR;
    var ret = this.copyout(it, value, err);
    if err then ioerror(err, "in buffer.copyout");
    return ret;
  }

  proc buffer.copyin( it:buffer_iterator, value, out error:syserr):buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      //writeln("iterator on way in");
      //debug_print_qbuffer_iter(it._bufit_internal);
      var end = it;
      var tmp = value;
      var sz = numBytes(value.type);
      //writeln("iterator copied to end");
      //debug_print_qbuffer_iter(end._bufit_internal);

      //writeln("sz is ", sz);
      this.advance(end, sz);
      //debug_print_qbuffer_iter(end._bufit_internal);
      error = qbuffer_copyin(this._buf_internal, it._bufit_internal, end._bufit_internal, tmp, sz);
      ret = end;
    }
    return ret;
  }
  proc buffer.copyin( it:buffer_iterator, value):buffer_iterator {
    var err:syserr = ENOERR;
    var ret = this.copyin(it, value, err);
    if err then ioerror(err, "in buffer.copyin");
    return ret;
  }
}

module BufferInternals {
  use SysBasic;
  use Error;

  extern type qbytes_ptr_t;
  extern type qbuffer_ptr_t;
  extern type qbuffer_iter_t;
  extern const QBYTES_PTR_NULL:qbytes_ptr_t;
  extern const QBUFFER_PTR_NULL:qbuffer_ptr_t;

  extern proc qbytes_retain(qb:qbytes_ptr_t);
  extern proc qbytes_release(qb:qbytes_ptr_t);
  extern proc qbytes_len(qb:qbytes_ptr_t):int(64);
  extern proc qbytes_data(qb:qbytes_ptr_t):c_void_ptr;

  extern proc qbytes_create_iobuf(out ret:qbytes_ptr_t):err_t;
  extern proc qbytes_create_calloc(out ret:qbytes_ptr_t, len:int(64)):err_t;

  extern proc qbuffer_iter_null():qbuffer_iter_t;

  /*
  extern proc qbuffer_init(ref buf:qbuffer_t):err_t;
  extern proc qbuffer_destroy(ref buf:qbuffer_t);
  */
  extern proc qbuffer_create(out buf:qbuffer_ptr_t):err_t;
  extern proc qbuffer_retain(buf:qbuffer_ptr_t);
  extern proc qbuffer_release(buf:qbuffer_ptr_t);

  extern proc qbuffer_append(buf:qbuffer_ptr_t, bytes:qbytes_ptr_t, skip_bytes:int(64), len_bytes:int(64)):err_t;
  extern proc qbuffer_append_buffer(buf:qbuffer_ptr_t, src:qbuffer_ptr_t, src_start:qbuffer_iter_t, src_end:qbuffer_iter_t):err_t;
  extern proc qbuffer_prepend(buf:qbuffer_ptr_t, bytes:qbytes_ptr_t, skip_bytes:int(64), len_bytes:int(64)):err_t;
  extern proc qbuffer_flatten(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, out bytes_out):err_t;
  //extern proc qbuffer_clone(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, out buf_out:qbuffer_t):err_t;
  extern proc qbuffer_copyout(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, ref x, size):err_t;
  extern proc qbuffer_copyin(buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t, ref x, size):err_t;

  extern proc qbuffer_begin(buf:qbuffer_ptr_t):qbuffer_iter_t;
  extern proc qbuffer_end(buf:qbuffer_ptr_t):qbuffer_iter_t;
  extern proc qbuffer_iter_next_part(buf:qbuffer_ptr_t, ref it:qbuffer_iter_t);
  extern proc qbuffer_iter_prev_part(buf:qbuffer_ptr_t, ref it:qbuffer_iter_t);
  extern proc qbuffer_iter_advance(buf:qbuffer_ptr_t, ref it:qbuffer_iter_t, amt:int(64));

  extern proc qbuffer_iter_get(it: qbuffer_iter_t, end:qbuffer_iter_t,
                                out bytes_out:qbytes_ptr_t,
                                out skip_out:int(64),
                                out len_out:int(64));
  extern proc qbuffer_iter_num_bytes(start:qbuffer_iter_t, end:qbuffer_iter_t):int(64);

  extern proc qbuffer_len(buf:qbuffer_ptr_t):int(64);

  extern proc debug_print_qbuffer_iter(/*const*/ ref it:qbuffer_iter_t);

  extern proc qbuffer_start_offset(buf:qbuffer_ptr_t):int(64);
  extern proc qbuffer_end_offset(buf:qbuffer_ptr_t):int(64);
  extern proc qbuffer_reposition(buf:qbuffer_ptr_t, new_offset_start:int(64));

  extern proc bulk_get_bytes(src_locale:int, src_addr:qbytes_ptr_t):qbytes_ptr_t;

  extern proc bulk_put_buffer(dst_locale:int, dst_addr:c_void_ptr, dst_len:int(64), buf:qbuffer_ptr_t, start:qbuffer_iter_t, end:qbuffer_iter_t):err_t;

}

module Buffers {
  use BufferInternals; // TODO -- non-exporting use

  // Now define the Chapel types using the originals..
  record bytes {
    var home: locale;
    var _bytes_internal:qbytes_ptr_t;
  }

  record buffer_iterator {
    var home: locale;
    var _bufit_internal:qbuffer_iter_t;
  }
  record buffer_range {
    var start:buffer_iterator;
    var end:buffer_iterator;
  }

  proc buffer_range.len:int(64) {
    var ret:int(64);
    on this.start.home {
      ret = qbuffer_iter_num_bytes(this.start._bufit_internal,
                                   this.end._bufit_internal);
    }
    return ret;
  }

  // FUTURE -- we could create const buffers
  // with a boolean type param
  record buffer {
    var home:locale;
    var _buf_internal:qbuffer_ptr_t;
  }

  // bytes methods.
  proc bytes.bytes() {
    this._bytes_internal = QBYTES_PTR_NULL;
    this.home = here;
  }
  proc bytes.bytes(len:int(64), out error:syserr) {
    error = qbytes_create_calloc(this._bytes_internal, len);
    // The buffer is "retained" internally on creation, but only on success.
    this.home = here;
  }
  proc bytes.bytes(len:int(64)) {
    var error:syserr = ENOERR;
    error = qbytes_create_calloc(this._bytes_internal, len);
    if error then ioerror(error, "in bytes constructor");
    // The buffer is retained internally on construction, but only on success.
    this.home = here;
  }

  proc create_iobuf(out error:syserr):bytes {
    var ret:bytes;
    error = qbytes_create_iobuf(ret._bytes_internal);
    // The buffer is "retained" internally on creation, but only on success.
    ret.home = here;
    return ret;
  }
  proc create_iobuf():bytes {
    var err:syserr = ENOERR;
    var ret = create_iobuf(err);
    if err then ioerror(err, "in create_iobuf");
    return ret;
  }

  // TODO -- shouldn't have to write this this way!
  pragma "init copy fn"
  proc chpl__initCopy(x: bytes) {
    if x.home == here {
      qbytes_retain(x._bytes_internal);
      return x;
    } else {
      var ret:bytes;
      ret.home = here;
      writeln("Bulk moving bytes");
      // The initial ref count is 1, so no need to call qbytes_retain here.
      ret._bytes_internal = bulk_get_bytes(x.home.id, x._bytes_internal);
      return ret;
    }
  }

  proc =(ref ret:bytes, x:bytes) {
    // retain -- release
    if( x.home == here ) {
      on x.home {
        qbytes_retain(x._bytes_internal);
      }

      on ret.home {
        qbytes_release(ret._bytes_internal);
      }

      ret.home = x.home;
      ret._bytes_internal = x._bytes_internal;
    } else {
      on ret.home {
        qbytes_release(ret._bytes_internal);
      }
      ret.home = here;
      writeln("Bulk moving bytes");
      ret._bytes_internal = bulk_get_bytes(x.home.id, x._bytes_internal);
      // On return from bulk_get_bytes, the ref count in ret._bytes_internal
      // should be 1.
      // Note that the error case is not handled (bulk_get_bytes must succeed).
    }
  }

  proc bytes.~bytes() {
    on this.home {
      qbytes_release(this._bytes_internal);
      this._bytes_internal = QBYTES_PTR_NULL;
    }
  }
  proc bytes.len:int(64) {
    var ret:int(64);
    on home {
     ret = qbytes_len(this._bytes_internal);
    }
    return ret;
  }

  // buffer iterator
  proc buffer_iterator.buffer_iterator() {
    this.home = here;
    this._bufit_internal = qbuffer_iter_null();
  }
  /* don't believe this is necessary...
  pragma "init copy fn"
  proc chpl__initCopy(x: buffer_iterator) {
    return x;
  }
  */

  // buffer methods.
  proc buffer.buffer(out error:syserr) {
    this.home = here;
    error = qbuffer_create(this._buf_internal);
  }
  proc buffer.buffer() {
    var error:syserr = ENOERR;
    this.home = here;
    error = qbuffer_create(this._buf_internal);
    if error then ioerror(error, "in buffer constructor");
  }

  // buffer.flatten is special in that it creates a
  // bytes object wherever it is called (and communicates it
  // intelligently). It is use in buffer's initCopy.
  proc buffer.flatten(range:buffer_range, out error:syserr):bytes {
    var ret:bytes;
    if this.home == here {
      error = qbuffer_flatten(this._buf_internal, range.start._bufit_internal, range.end._bufit_internal, ret);
    } else {
      var dst_locale = here;
      var dst_len:int(64) = range.len;
      ret = new bytes(dst_len);
      var dst_addr = qbytes_data(ret._bytes_internal);
      on this.home {
        // Copy the buffer to the bytes...
        error = bulk_put_buffer(dst_locale.id, dst_addr, dst_len,
                                this._buf_internal,
                                range.start._bufit_internal,
                                range.end._bufit_internal);
      }
    }
    return ret;
  }

  // TODO -- shouldn't have to write this this way!
  pragma "init copy fn"
  proc chpl__initCopy(x: buffer) {
    if x.home == here {
      qbuffer_retain(x._buf_internal);
      return x;
    } else {
      var ret:buffer;
      var start_offset:int(64);
      var end_offset:int(64);

      on x.home {
        start_offset = qbuffer_start_offset(x._buf_internal);
        end_offset = qbuffer_end_offset(x._buf_internal);
      }

      var b = new bytes(end_offset - start_offset);
      var len = qbytes_len(b._bytes_internal);
      var ptr = qbytes_data(b._bytes_internal);

      // Now create a local buffer with the right data
      // in it starting at the right position.
      qbuffer_reposition(ret._buf_internal, start_offset);

      var there_uid = here.id;

      on x.home {
        var err:syserr = ENOERR;
        err = bulk_put_buffer(there_uid, ptr, len, x._buf_internal,
                              qbuffer_begin(x._buf_internal),
                              qbuffer_end(x._buf_internal));
        if err then ioerror(err, "in buffer init copy");
      }

      ret.append(b);
      return ret;
    }
  }

  proc =(ref ret:buffer, x:buffer) {
    ret.home = here;
    // retain -- release
    if( x.home == ret.home ) {
      on x.home {
        qbuffer_retain(x._buf_internal);
      }

      on ret.home {
        qbuffer_release(ret._buf_internal);
      }

      ret.home = x.home;
      ret._buf_internal = x._buf_internal;
    } else {
      on ret.home {
        qbuffer_release(ret._buf_internal);
      }

      var start_offset:int(64);
      var end_offset:int(64);

      on x.home {
        start_offset = qbuffer_start_offset(x._buf_internal);
        end_offset = qbuffer_end_offset(x._buf_internal);
      }

      var b = new bytes(end_offset - start_offset);
      var len = qbytes_len(b._bytes_internal);
      var ptr = qbytes_data(b._bytes_internal);

      // Now create a local buffer with the right data
      // in it starting at the right position.
      qbuffer_reposition(ret._buf_internal, start_offset);

      var there_uid = here.id;

      on x.home {
        var err:syserr = ENOERR;
        err = bulk_put_buffer(there_uid, ptr, len, x._buf_internal,
                              qbuffer_begin(x._buf_internal),
                              qbuffer_end(x._buf_internal));
        if err then ioerror(err, "in buffer assignment");
      }

      ret.append(b);
    }
  }

  proc buffer.~buffer() {
    on this.home {
      qbuffer_release(this._buf_internal);
      this._buf_internal = QBUFFER_PTR_NULL;
    }
  }

  proc buffer.len {
    var ret:int(64);
    on this.home {
      ret = qbuffer_len(this._buf_internal);
    }
    return ret;
  }

  proc buffer.append(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len, out error:syserr) {
    on this.home {
      error = qbuffer_append(this._buf_internal, b._bytes_internal, skip_bytes, len_bytes);
    }
  }
  proc buffer.append(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len) {
    var err:syserr = ENOERR;
    this.append(b, skip_bytes, len_bytes, err);
    if err then ioerror(err, "in buffer.append");
  }

  proc buffer.append(buf:buffer, part:buffer_range = buf.all(), out error:syserr) {
    on this.home {
      error = qbuffer_append_buffer(this._buf_internal, buf._buf_internal, part.start._bufit_internal, part.end._bufit_internal);
    }
  }
  proc buffer.append(buf:buffer, part:buffer_range = buf.all()) {
    var err:syserr = ENOERR;
    this.append(buf, part, err);
    if err then ioerror(err, "in buffer.append");
  }

  proc buffer.prepend(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len, out error:syserr) {
    on this.home {
      qbuffer_prepend(this._buf_internal, b._bytes_internal, skip_bytes, len_bytes);
    }
  }
  proc buffer.prepend(b:bytes, skip_bytes:int(64) = 0, len_bytes:int(64) = b.len) {
    var err:syserr = ENOERR;
    this.prepend(b, skip_bytes, len_bytes, err);
    if err then ioerror(err, "in buffer.prepend");
  }

  proc buffer.start():buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      ret._bufit_internal = qbuffer_begin(this._buf_internal);
    }
    return ret;
  }
  proc buffer.end():buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      ret._bufit_internal = qbuffer_end(this._buf_internal);
    }
    return ret;
  }
  proc buffer.all():buffer_range {
    return new buffer_range(this.start(), this.end());
  }

  proc buffer.next_part(inout it:buffer_iterator) {
    on this.home {
      var x = it;
      qbuffer_iter_next_part(this._buf_internal, x._bufit_internal);
      it = x;
    }
  }
  proc buffer.prev_part(inout it:buffer_iterator) {
    on this.home {
      var x = it;
      qbuffer_iter_prev_part(this._buf_internal, x._bufit_internal);
      it = x;
    }
  }
  proc buffer.advance(inout it:buffer_iterator, amount:int(64)) {
    on this.home {
      var x = it;
      qbuffer_iter_advance(this._buf_internal, x._bufit_internal, amount);
      it = x;
    }
  }

  proc buffer_iterator.debug_print()
  {
    on this.home {
      debug_print_qbuffer_iter(this._bufit_internal);
    }
  }

  /*
  proc buffer.clone(range:buffer_range, err:ErrorHandler = nil):buffer {
    on this.home {
      var ret:buffer;
      seterr(qbuffer_clone(this._buf_internal, range.start, range.end, ret._buf_internal));
      return ret;
    }
  }*/

  // methods to read/write basic types.
  proc buffer.copyout(it:buffer_iterator, out value, out error:syserr):buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      var end:buffer_iterator = it;
      var tmp:value.type;
      var sz = numBytes(value.type);
      this.advance(end, sz);
      error = qbuffer_copyout(this._buf_internal, it._bufit_internal, end._bufit_internal, tmp, sz);
      value = tmp;
      ret = end;
    }
    return ret;
  }
  proc buffer.copyout(it:buffer_iterator, out value):buffer_iterator {
    var err:syserr = ENOERR;
    this.copyout(it, value, err);
    if err then ioerror(err, "in buffer.copyout");
  }

  proc buffer.copyin( it:buffer_iterator, value, out error:syserr):buffer_iterator {
    var ret:buffer_iterator;
    ret.home = this.home;
    on this.home {
      //writeln("iterator on way in");
      //debug_print_qbuffer_iter(it._bufit_internal);
      var end = it;
      var tmp = value;
      var sz = numBytes(value.type);
      //writeln("iterator copied to end");
      //debug_print_qbuffer_iter(end._bufit_internal);

      //writeln("sz is ", sz);
      this.advance(end, sz);
      //debug_print_qbuffer_iter(end._bufit_internal);
      error = qbuffer_copyin(this._buf_internal, it._bufit_internal, end._bufit_internal, tmp, sz);
      ret = end;
    }
    return ret;
  }
  proc buffer.copyin( it:buffer_iterator, value):buffer_iterator {
    var err:syserr = ENOERR;
    this.copyin(it, value, err);
    if err then ioerror(err, "in buffer.copyin");
  }
}


Switch to nested foralls

Switch to nested foralls

use Histogram;
use Random;
use Time;

// Test flags
config const isTest=false;
config const isPerf=false;
config const doBrute=false;

// Performance test parameters
config const nParticles=10000;

// Input/Output filenames
config const fn1 = "test.dat";
config const fn2 = "test.dat";
config const pairfn = "test-DD.dat";

// Dimension constanst
param NDIM=3;
param NEXT=2; // w, r2
param NTOT=NDIM+NEXT;
const W=NDIM;
const R2=NDIM+1;
const ParticleAttrib = 0.. #NTOT;
const DimSpace = 0.. #NDIM;

// Tree parameters
config const minpart=500;

// Histogram parameters
config const smax=200.0;
config const smax2=smax**2;
config const nmubins=5;
config const nsbins=5;

// Global variables
var gtime1 : Timer;

proc main() {
  doPairs();
}

class Particle3D {
  var npart : int = 0;
  var Darr : domain(2);
  var Dpart : domain(1);
  var arr : [Darr] real;

  var _tmp : [Dpart] real;
  var _n1, _ndx : [Dpart] int;

  proc Particle3D(npart1 : int, random : bool = false) {
    npart = npart1;
    Darr = {ParticleAttrib, 0.. #npart};
    Dpart = {0.. #npart};
    if random {
      var rng = new RandomStream();
      var x, y, z : real;
      for ii in Dpart {
        x = rng.getNext()*1000.0; y = rng.getNext()*1000.0; z = rng.getNext()*1000.0;
        arr[0,ii] = x; arr[1, ii] = y; arr[2, ii] = z;
        arr[3,ii] = 1.0;
        arr[4,ii] = x**2 + y**2 + z**2;
      }
      delete rng;
    }
  }

  proc this(ii : int) : NTOT*real {
    // ??? Hardcoded hack!
    return (arr[0,ii],arr[1,ii],arr[2,ii],arr[3,ii],arr[4,ii]);
  }

  // Reorder all the arrays according to _ndx (we assume you have set this)
  proc reorder(idom : domain(1)) {

    for kk in ParticleAttrib {
      forall ii in idom {
        _tmp[_ndx[ii]] = arr[kk,ii];
      }
      forall ii in idom {
        arr[kk,ii] = _tmp[ii];
      }
    }
  }

  // Shuffles the elements of the array...
  // Note that the random number generator cannot cover the space
  // of all permutations, but that's fine for what we need.
  proc shuffle() {
    // Fill in _ndx
    forall ii in Dpart {
      _ndx[ii] = ii;
    }

    // Set the random number generator
    var rng = new RandomStream(41);
    var jj : int;
    for ii in 0..(npart-2) {
      jj = (rng.getNext()*(npart-ii)):int + ii;
      _ndx[jj] <=> _ndx[ii];
    }
    delete rng;

    reorder(Dpart);
  }

}

proc countLines(fn : string) : int {
  var ff = open(fn, iomode.r);
  var ipart = 0;
  for iff in ff.lines() do ipart +=1;
  ff.close();
  return ipart;
}

proc readFile(fn : string) : Particle3D  {
  var npart = countLines(fn);
  var pp = new Particle3D(npart);

  var ff = openreader(fn);
  var ipart = 0;
  var x,y,z,w,r2 : real;
  while (ff.read(x,y,z,w)) {
    r2 = x**2 + y**2 + z**2;
    pp.arr[..,ipart] = [x,y,z,w,r2];
    ipart += 1;
  }
  ff.close();

  return pp;
}

proc splitOn(pp : Particle3D, idom : domain(1), splitDim : int, xsplit : real) : int {
  // Setup for a prefix scan
  forall ii in idom {
    if (pp.arr[splitDim,ii] < xsplit) {
      pp._n1[ii] = 1;
    } else {
      pp._n1[ii] = 0;
    }
  }

  var lnpart : int = 0;
  for ii in idom do lnpart += pp._n1[ii];
  var li = idom.low;
  var ri = li+lnpart;
  for ii in idom {
    if (pp._n1[ii] == 1) {
      pp._ndx[ii] = li;
      li += 1;
    } else {
      pp._ndx[ii] = ri;
      ri += 1;
    }
  }

  pp.reorder(idom);

  return lnpart;
}

class KDNode {
  var lo, hi,npart,id : int;
  var dom : domain(1);
  var xcen : [DimSpace]real;
  var rcell : real;
  var left, right : KDNode;

  proc isLeaf() : bool {
    return (left==nil) && (right==nil);
  }

  proc ~KDNode() {
    if left then delete left;
    if right then delete right;
  }

}

proc BuildTree(pp : Particle3D, lo : int, hi : int, id : int) : KDNode  {
  var me : KDNode = new KDNode();
  me.lo = lo;
  me.hi = hi;
  me.dom = {lo..hi};
  me.id = id;
  me.npart = (me.hi-me.lo)+1;

  //  work out xcen and vantage point radius
  var pmin, pmax : [DimSpace] real;
  for idim in DimSpace {
    pmin[idim] = min reduce pp.arr[idim,lo..hi];
    pmax[idim] = max reduce pp.arr[idim,lo..hi];
  }
  me.xcen = (pmax+pmin)/2.0;
  me.rcell = 0.0;
  var r1 : real;
  for ii in lo..hi {
    r1 = 0.0;
    for idim in DimSpace {
      r1 += (pp.arr[idim,ii]-me.xcen[idim])**2;
    }
    r1 = sqrt(r1);
    if (r1 > me.rcell) then me.rcell = r1;
  }

  // Continue to split
  if (me.npart <= minpart) then return me; // Don't split further.

  // Find dimension to split on
  var dx : [DimSpace] real;
  dx = pmax - pmin;
  var splitDim = 0;
  for idim in DimSpace {
    if (dx[idim] > dx[splitDim]) then splitDim=idim;
  }

  // Split
  var lnpart = splitOn(pp,me.dom,splitDim, me.xcen[splitDim]);
  me.left  = BuildTree(pp, lo, lo+lnpart-1, 2*id+1);
  me.right = BuildTree(pp, lo+lnpart, hi,2*id+2);
  return me;
}

proc TreeAccumulate(hh : UniformBins, p1, p2 : Particle3D, node1, node2 : KDNode) {
  // Compute the distance between node1 and node2
  var rr = sqrt (+ reduce(node1.xcen - node2.xcen)**2);
  var rmin = rr - (node1.rcell+node2.rcell);

  // If distance is greater than all cases
  if (rmin > smax) then return;

  // If both nodes are leaves
  if (node1.isLeaf() & node2.isLeaf()) {
    begin smuAccumulate(hh, p1, p2, node1.dom, node2.dom,1);
    return;
  }

  // If one node is a leaf
  if (node1.isLeaf()) {
    TreeAccumulate(hh, p1, p2, node1, node2.left);
    TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }
  if (node2.isLeaf()) {
    TreeAccumulate(hh, p1, p2, node1.left, node2);
    TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  }

  // Split the larger case;
  if (node1.npart > node2.npart) {
    TreeAccumulate(hh, p1, p2, node1.left, node2);
    TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  } else {
    TreeAccumulate(hh, p1, p2, node1, node2.left);
    TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }

}

// The basic pair counter
proc smuAccumulate(hh : UniformBins, p1,p2 : Particle3D, d1,d2 : domain(1), scale : real) {
  forall ii in d1 { // Loop over first set of particles

    var x1,y1,z1,w1,r2 : real;
    (x1,y1,z1,w1,r2) = p1[ii];

    forall jj in d2 { // Second set of particles
      var sl, s2, l1, s1, l2, mu, wprod : real;
      mu=2*(p2.arr[0,jj]*x1 + p2.arr[1,jj]*y1 + p2.arr[2,jj]*z1);
      sl = r2 - p2.arr[R2,jj];
      l1 = r2 + p2.arr[R2,jj];
      s2 = l1 - mu;
      l2 = l1 + mu;
      if ((s2 < smax2) && (s2 > 1.0e-20)) {
        wprod = scale * w1 * p2.arr[W,jj];
        s1 = sqrt(s2);
        mu = sl/(s1*sqrt(l2));
        if (mu < 0) then mu = -mu;

        hh.add((s1,mu),wprod);
      }
    }
  }
}

proc doPairs() {
  var tt : Timer;

  // Read in the file
  tt.clear(); tt.start();
  var pp1, pp2 : Particle3D;
  if isPerf {
    pp1 = new Particle3D(nParticles, true);
    pp2 = new Particle3D(nParticles, true);
  } else {
    pp1 = readFile(fn1);
    pp2 = readFile(fn2);
    if !isTest {
      writef("Read in %i lines from file %s \n", pp1.npart, fn1);
      writef("Read in %i lines from file %s \n", pp2.npart, fn2);
    }
  }
  tt.stop();
  if !isTest {
    writef("Time to read : %r \n", tt.elapsed());
  }

  // Shuffle the particles -- this is not used here, but will be usef for the multilocale version
  tt.clear(); tt.start();
  pp1.shuffle();
  pp2.shuffle();
  tt.stop();
  if !isTest then writef("Time to shuffle : %r \n",tt.elapsed());

  // Build the tree
  tt.clear(); tt.start(); gtime1.clear();
  var root1 = BuildTree(pp1, 0, pp1.npart-1, 0);
  var root2 = BuildTree(pp2, 0, pp2.npart-1, 0);
  tt.stop();
  if !isTest then writef("Time to build trees : %r \n", tt.elapsed());
  //writef("Time in splitOn : %r \n", gtime1.elapsed());

  // Set up the histogram
  var hh = new UniformBins(2,(nsbins,nmubins), ((0.0,smax),(0.0,1.0+1.e-10)));

  // Do the paircounts with a tree
  hh.reset();
  tt.clear(); tt.start();
  sync TreeAccumulate(hh, pp1,pp2, root1, root2);
  tt.stop();
  if (!isTest) {
    writef("Time to tree paircount : %r \n", tt.elapsed());
    if !isPerf {
      var ff = openwriter("%s.tree".format(pairfn));
      writeHist(ff,hh);
      ff.close();
    }
  } else {
    hh.set((0,0),0.0);
    writeHist(stdout,hh,"%20.5er ");
  }

  //
  // clean up
  //
  delete pp1;
  delete pp2;
  delete root1;
  delete root2;
  delete hh;
}

use Histogram;
use Random;
use Time;

// Test flags
config const isTest=false;
config const isPerf=false;
config const doBrute=false;

// Performance test parameters
config const nParticles=10000;

// Input/Output filenames
config const fn1 = "test.dat";
config const fn2 = "test.dat";
config const pairfn = "test-DD.dat";

// Dimension constanst
param NDIM=3;
param NEXT=2; // w, r2
param NTOT=NDIM+NEXT;
const W=NDIM;
const R2=NDIM+1;
const ParticleAttrib = 0.. #NTOT;
const DimSpace = 0.. #NDIM;

// Tree parameters
config const minpart=500;

// Histogram parameters
config const smax=200.0;
config const smax2=smax**2;
config const nmubins=5;
config const nsbins=5;

// Global variables
var gtime1 : Timer;

proc main() {
  doPairs();
}

class Particle3D {
  var npart : int = 0;
  var Darr : domain(2);
  var Dpart : domain(1);
  var arr : [Darr] real;

  var _tmp : [Dpart] real;
  var _n1, _ndx : [Dpart] int;

  proc Particle3D(npart1 : int, random : bool = false) {
    npart = npart1;
    Darr = {ParticleAttrib, 0.. #npart};
    Dpart = {0.. #npart};
    if random {
      var rng = new RandomStream();
      var x, y, z : real;
      for ii in Dpart {
        x = rng.getNext()*1000.0; y = rng.getNext()*1000.0; z = rng.getNext()*1000.0;
        arr[0,ii] = x; arr[1, ii] = y; arr[2, ii] = z;
        arr[3,ii] = 1.0;
        arr[4,ii] = x**2 + y**2 + z**2;
      }
      delete rng;
    }
  }

  proc this(ii : int) : NTOT*real {
    // ??? Hardcoded hack!
    return (arr[0,ii],arr[1,ii],arr[2,ii],arr[3,ii],arr[4,ii]);
  }

  // Reorder all the arrays according to _ndx (we assume you have set this)
  proc reorder(idom : domain(1)) {

    for kk in ParticleAttrib {
      forall ii in idom {
        _tmp[_ndx[ii]] = arr[kk,ii];
      }
      forall ii in idom {
        arr[kk,ii] = _tmp[ii];
      }
    }
  }

  // Shuffles the elements of the array...
  // Note that the random number generator cannot cover the space
  // of all permutations, but that's fine for what we need.
  proc shuffle() {
    // Fill in _ndx
    forall ii in Dpart {
      _ndx[ii] = ii;
    }

    // Set the random number generator
    var rng = new RandomStream(41);
    var jj : int;
    for ii in 0..(npart-2) {
      jj = (rng.getNext()*(npart-ii)):int + ii;
      _ndx[jj] <=> _ndx[ii];
    }
    delete rng;

    reorder(Dpart);
  }

}

proc countLines(fn : string) : int {
  var ff = open(fn, iomode.r);
  var ipart = 0;
  for iff in ff.lines() do ipart +=1;
  ff.close();
  return ipart;
}

proc readFile(fn : string) : Particle3D  {
  var npart = countLines(fn);
  var pp = new Particle3D(npart);

  var ff = openreader(fn);
  var ipart = 0;
  var x,y,z,w,r2 : real;
  while (ff.read(x,y,z,w)) {
    r2 = x**2 + y**2 + z**2;
    pp.arr[..,ipart] = [x,y,z,w,r2];
    ipart += 1;
  }
  ff.close();

  return pp;
}

proc splitOn(pp : Particle3D, idom : domain(1), splitDim : int, xsplit : real) : int {
  // Setup for a prefix scan
  forall ii in idom {
    if (pp.arr[splitDim,ii] < xsplit) {
      pp._n1[ii] = 1;
    } else {
      pp._n1[ii] = 0;
    }
  }

  var lnpart : int = 0;
  for ii in idom do lnpart += pp._n1[ii];
  var li = idom.low;
  var ri = li+lnpart;
  for ii in idom {
    if (pp._n1[ii] == 1) {
      pp._ndx[ii] = li;
      li += 1;
    } else {
      pp._ndx[ii] = ri;
      ri += 1;
    }
  }

  pp.reorder(idom);

  return lnpart;
}

class KDNode {
  var lo, hi,npart,id : int;
  var dom : domain(1);
  var xcen : [DimSpace]real;
  var rcell : real;
  var left, right : KDNode;

  proc isLeaf() : bool {
    return (left==nil) && (right==nil);
  }

  proc ~KDNode() {
    if left then delete left;
    if right then delete right;
  }

}

proc BuildTree(pp : Particle3D, lo : int, hi : int, id : int) : KDNode  {
  var me : KDNode = new KDNode();
  me.lo = lo;
  me.hi = hi;
  me.dom = {lo..hi};
  me.id = id;
  me.npart = (me.hi-me.lo)+1;

  //  work out xcen and vantage point radius
  var pmin, pmax : [DimSpace] real;
  for idim in DimSpace {
    pmin[idim] = min reduce pp.arr[idim,lo..hi];
    pmax[idim] = max reduce pp.arr[idim,lo..hi];
  }
  me.xcen = (pmax+pmin)/2.0;
  me.rcell = 0.0;
  var r1 : real;
  for ii in lo..hi {
    r1 = 0.0;
    for idim in DimSpace {
      r1 += (pp.arr[idim,ii]-me.xcen[idim])**2;
    }
    r1 = sqrt(r1);
    if (r1 > me.rcell) then me.rcell = r1;
  }

  // Continue to split
  if (me.npart <= minpart) then return me; // Don't split further.

  // Find dimension to split on
  var dx : [DimSpace] real;
  dx = pmax - pmin;
  var splitDim = 0;
  for idim in DimSpace {
    if (dx[idim] > dx[splitDim]) then splitDim=idim;
  }

  // Split
  var lnpart = splitOn(pp,me.dom,splitDim, me.xcen[splitDim]);
  me.left  = BuildTree(pp, lo, lo+lnpart-1, 2*id+1);
  me.right = BuildTree(pp, lo+lnpart, hi,2*id+2);
  return me;
}

proc TreeAccumulate(hh : UniformBins, p1, p2 : Particle3D, node1, node2 : KDNode) {
  // Compute the distance between node1 and node2
  var rr = sqrt (+ reduce(node1.xcen - node2.xcen)**2);
  var rmin = rr - (node1.rcell+node2.rcell);

  // If distance is greater than all cases
  if (rmin > smax) then return;

  // If both nodes are leaves
  if (node1.isLeaf() & node2.isLeaf()) {
    begin smuAccumulate(hh, p1, p2, node1.dom, node2.dom,1);
    return;
  }

  // If one node is a leaf
  if (node1.isLeaf()) {
    TreeAccumulate(hh, p1, p2, node1, node2.left);
    TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }
  if (node2.isLeaf()) {
    TreeAccumulate(hh, p1, p2, node1.left, node2);
    TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  }

  // Split the larger case;
  if (node1.npart > node2.npart) {
    TreeAccumulate(hh, p1, p2, node1.left, node2);
    TreeAccumulate(hh, p1, p2, node1.right, node2);
    return;
  } else {
    TreeAccumulate(hh, p1, p2, node1, node2.left);
    TreeAccumulate(hh, p1, p2, node1, node2.right);
    return;
  }

}

// The basic pair counter
proc smuAccumulate(hh : UniformBins, p1,p2 : Particle3D, d1,d2 : domain(1), scale : real) {
  forall ii in d1 { // Loop over first set of particles

    var x1,y1,z1,w1,r2 : real;
    var sl, s2, l1, s1, l2, mu, wprod : real;
    (x1,y1,z1,w1,r2) = p1[ii];

    for jj in d2 { // Second set of particles
      mu=2*(p2.arr[0,jj]*x1 + p2.arr[1,jj]*y1 + p2.arr[2,jj]*z1);
      sl = r2 - p2.arr[R2,jj];
      l1 = r2 + p2.arr[R2,jj];
      s2 = l1 - mu;
      l2 = l1 + mu;
      if ((s2 >= smax2) || (s2 < 1.0e-20)) then continue;
      wprod = scale * w1 * p2.arr[W,jj];
      s1 = sqrt(s2);
      mu = sl/(s1*sqrt(l2));
      if (mu < 0) then mu = -mu;

      hh.add((s1,mu),wprod);
    }
  }
}

proc doPairs() {
  var tt : Timer;

  // Read in the file
  tt.clear(); tt.start();
  var pp1, pp2 : Particle3D;
  if isPerf {
    pp1 = new Particle3D(nParticles, true);
    pp2 = new Particle3D(nParticles, true);
  } else {
    pp1 = readFile(fn1);
    pp2 = readFile(fn2);
    if !isTest {
      writef("Read in %i lines from file %s \n", pp1.npart, fn1);
      writef("Read in %i lines from file %s \n", pp2.npart, fn2);
    }
  }
  tt.stop();
  if !isTest {
    writef("Time to read : %r \n", tt.elapsed());
  }

  // Shuffle the particles -- this is not used here, but will be usef for the multilocale version
  tt.clear(); tt.start();
  pp1.shuffle();
  pp2.shuffle();
  tt.stop();
  if !isTest then writef("Time to shuffle : %r \n",tt.elapsed());

  // Build the tree
  tt.clear(); tt.start(); gtime1.clear();
  var root1 = BuildTree(pp1, 0, pp1.npart-1, 0);
  var root2 = BuildTree(pp2, 0, pp2.npart-1, 0);
  tt.stop();
  if !isTest then writef("Time to build trees : %r \n", tt.elapsed());
  //writef("Time in splitOn : %r \n", gtime1.elapsed());

  // Set up the histogram
  var hh = new UniformBins(2,(nsbins,nmubins), ((0.0,smax),(0.0,1.0+1.e-10)));

  // Do the paircounts with a tree
  hh.reset();
  tt.clear(); tt.start();
  sync TreeAccumulate(hh, pp1,pp2, root1, root2);
  tt.stop();
  if (!isTest) {
    writef("Time to tree paircount : %r \n", tt.elapsed());
    if !isPerf {
      var ff = openwriter("%s.tree".format(pairfn));
      writeHist(ff,hh);
      ff.close();
    }
  } else {
    hh.set((0,0),0.0);
    writeHist(stdout,hh,"%20.5er ");
  }

  //
  // clean up
  //
  delete pp1;
  delete pp2;
  delete root1;
  delete root2;
  delete hh;
}


Fixed array test to use square brackets on domain.  This doesn't work
because of the line "A(i) = i" though it looks to me like a coercion
wrapper is built correctly.

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@4914 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

config var n: integer = 4;

var D: domain(1) = [1..n];

var A: [D] float;

forall i in D {
  A(i) = i;
}

function foo(X) {
  writeln("X is: ", X);
}

foo(A);

config var n: integer = 4;

var D: domain(1) = (1..n);

var A: [D] float;

forall i in D {
  A(i) = i;
}

function foo(X) {
  writeln("X is: ", X);
}

foo(A);


Minor bugfix

Minor bugfix

// Wrapper for FFTW version 3.
module FFTW {

  use SysCTypes;

  // Define the various planner flags
  // See Sec. 4.3.2 of FFTW manual "Planner Flags"
  extern const FFTW_FORWARD : c_int;
  extern const FFTW_BACKWARD : c_int;
  extern const FFTW_ESTIMATE : c_uint;
  extern const FFTW_MEASURE : c_uint;
  extern const FFTW_PATIENT : c_uint;
  extern const FFTW_EXHAUSTIVE : c_uint;
  extern const FFTW_WISDOM_ONLY : c_uint;
  extern const FFTW_DESTROY_INPUT : c_uint;
  extern const FFTW_PRESERVE_INPUT : c_uint;
  extern const FFTW_UNALIGNED : c_uint;

  // Define types
  // NOTE : Ideally, this should be mapped to complex(128), but I seem to run into
  // casting issues in the C code. This is a simple workaround for now - although ugly.
  // Also note that if one used complex types, then one would need to include complex.h at the command line
  extern type fftw_complex = 2*real(64); // 4.1.1
  extern type fftw_plan; // opaque type

  // Planner functions
  // Complex : 4.3.1
  // NOTE : We pass in arrays using ref
  proc plan_dft(in1 : [] fftw_complex, out1 : [] fftw_complex, sign : c_int, flags :c_uint) : fftw_plan
  {
    param rank = in1.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = in1.domain.dim(i).size: c_int;

    extern proc fftw_plan_dft(rank: c_int,
        n,  // BLC: having trouble being specific
        in1: [] fftw_complex,
        out1: [] fftw_complex,
        sign : c_int, flags : c_uint) : fftw_plan;

    return fftw_plan_dft(rank, dims, in1, out1, sign, flags);
  }

  // Real-to-complex and complex-to-real planning routines
  // There are two cases that we treat independently here : in-place and out-of-place transforms

  // Since the calls to FFTW are the same, pull the extern declarations out.
  // TODO : This should be cleaned up further and made consistent across the file
  extern proc fftw_plan_dft_r2c(rank: c_int,
      n,  // BLC: having trouble being specific
      in1: [],
      out1: [],
      flags : c_uint) : fftw_plan;
  extern proc fftw_plan_dft_c2r(rank: c_int,
      n,  // BLC: having trouble being specific
      in1: [],
      out1: [],
      flags : c_uint) : fftw_plan;

  // Out-of-place routines
  proc plan_dft_r2c(in1 : [] real(64), out1 : [] fftw_complex, flags :c_uint) : fftw_plan
  {
    param rank = in1.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = in1.domain.dim(i).size: c_int;

    return fftw_plan_dft_r2c(rank, dims, in1, out1, flags);
  }
  proc plan_dft_c2r(in1 : [] fftw_complex, out1 : [] real(64),  flags :c_uint) : fftw_plan
  {
    param rank = out1.rank: c_int; // The dimensions are that of the real array

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = out1.domain.dim(i).size: c_int;

    return fftw_plan_dft_c2r(rank, dims, in1, out1, flags);
  }

  // In-place routines, note that these take in the true leading dimension
  // TODO : We should put in checks to see that the sizes are consistent
  // TODO : We should check on types...
  proc plan_dft_r2c(realDom : domain, in1 : [] ?t, flags : c_uint) : fftw_plan
  {
    param rank = realDom.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = realDom.dim(i).size: c_int;

    return fftw_plan_dft_r2c(rank, dims, in1, in1, flags);
  }
  proc plan_dft_c2r(realDom : domain, in1: [] ?t, flags : c_uint) : fftw_plan
  {
    param rank = realDom.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = realDom.dim(i).size: c_int;

    return fftw_plan_dft_c2r(rank, dims, in1, in1, flags);
  }

  // Using plans
  proc execute(const plan : fftw_plan) {
    extern proc fftw_execute(const plan : fftw_plan);
    fftw_execute(plan);
  }
  proc destroy_plan(plan : fftw_plan) {
    extern proc fftw_destroy_plan(plan : fftw_plan);
    fftw_destroy_plan(plan);
  }
  proc cleanup() {
    extern proc fftw_cleanup();
    fftw_cleanup();
  }

  // Multi-threaded setup routines below
  proc init_threads() : c_int {
    extern proc fftw_init_threads() : c_int;
    return fftw_init_threads();
  }
  proc cleanup_threads() {
    extern proc fftw_cleanup_threads();
    fftw_cleanup_threads();
  }
  proc plan_with_nthreads(n : c_int) {
    extern proc fftw_plan_with_nthreads(n : c_int);
    fftw_plan_with_nthreads(n);
  }

  // Utilities -- not in FFTW
  proc abs((r,c) : fftw_complex) : real(64) {
    return sqrt(r*r + c*c);
  }

  proc re((r,c) : fftw_complex) : real(64) {
    return r;
  }
  proc im((r,c) : fftw_complex) : real(64) {
    return c;
  }

}

// Wrapper for FFTW version 3.
module FFTW {

  use SysCTypes;

  // Define the various planner flags
  // See Sec. 4.3.2 of FFTW manual "Planner Flags"
  extern const FFTW_FORWARD : c_int;
  extern const FFTW_BACKWARD : c_int;
  extern const FFTW_ESTIMATE : c_uint;
  extern const FFTW_MEASURE : c_uint;
  extern const FFTW_PATIENT : c_uint;
  extern const FFTW_EXHAUSTIVE : c_uint;
  extern const FFTW_WISDOM_ONLY : c_uint;
  extern const FFTW_DESTROY_INPUT : c_uint;
  extern const FFTW_PRESERVE_INPUT : c_uint;
  extern const FFTW_UNALIGNED : c_uint;

  // Define types
  // NOTE : Ideally, this should be mapped to complex(128), but I seem to run into
  // casting issues in the C code. This is a simple workaround for now - although ugly.
  // Also note that if one used complex types, then one would need to include complex.h at the command line
  extern type fftw_complex = 2*real(64); // 4.1.1
  extern type fftw_plan; // opaque type

  // Planner functions
  // Complex : 4.3.1
  // NOTE : We pass in arrays using ref
  proc plan_dft(in1 : [] fftw_complex, out1 : [] fftw_complex, sign : c_int, flags :c_uint) : fftw_plan
  {
    param rank = in1.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = in1.domain.dim(i).size: c_int;

    extern proc fftw_plan_dft(rank: c_int,
        n,  // BLC: having trouble being specific
        in1: [] fftw_complex,
        out1: [] fftw_complex,
        sign : c_int, flags : c_uint) : fftw_plan;

    return fftw_plan_dft(rank, dims, in1, out1, sign, flags);
  }

  // Real-to-complex and complex-to-real planning routines
  // There are two cases that we treat independently here : in-place and out-of-place transforms

  // Since the calls to FFTW are the same, pull the extern declarations out.
  // TODO : This should be cleaned up further and made consistent across the file
  extern proc fftw_plan_dft_r2c(rank: c_int,
      n,  // BLC: having trouble being specific
      in1: [],
      out1: [],
      flags : c_uint) : fftw_plan;
  extern proc fftw_plan_dft_c2r(rank: c_int,
      n,  // BLC: having trouble being specific
      in1: [],
      out1: [],
      flags : c_uint) : fftw_plan;

  // Out-of-place routines
  proc plan_dft_r2c(in1 : [] real(64), out1 : [] fftw_complex, flags :c_uint) : fftw_plan
  {
    param rank = in1.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = in1.domain.dim(i).size: c_int;

    return fftw_plan_dft_r2c(rank, dims, in1, out1, flags);
  }
  proc plan_dft_c2r(in1 : [] fftw_complex, out1 : [] real(64),  flags :c_uint) : fftw_plan
  {
    param rank = out1.rank: c_int; // The dimensions are that of the real array

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = out1.domain.dim(i).size: c_int;

    return fftw_plan_dft_c2r(rank, dims, in1, out1, flags);
  }

  // In-place routines, note that these take in the true leading dimension
  // TODO : We should put in checks to see that the sizes are consistent
  // TODO : We should check on types...
  proc plan_dft_r2c(realDom : domain, in1 : [] ?t, flags : c_uint) : fftw_plan
  {
    param rank = realDom.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = realDom.dim(i).size: c_int;

    return fftw_plan_dft_r2c(rank, dims, in1, in1, flags);
  }
  proc plan_dft_c2r(realDom : domain, in1: [] ?t, flags : c_uint) : fftw_plan
  {
    param rank = realDom.rank: c_int;

    var dims: rank*c_int;
    for param i in 1..rank do
      dims(i) = realDom.dim(i).size: c_int;

    return fftw_plan_dft_c2r(rank, dims, in1, in1, flags);
  }

  // Using plans
  proc execute(const plan : fftw_plan) {
    extern proc fftw_execute(const plan : fftw_plan);
    fftw_execute(plan);
  }
  proc destroy_plan(plan : fftw_plan) {
    extern proc fftw_destroy_plan(plan : fftw_plan);
    fftw_destroy_plan(plan);
  }
  proc cleanup() {
    extern proc fftw_cleanup();
    cleanup();
  }

  // Multi-threaded setup routines below
  proc init_threads() : c_int {
    extern proc fftw_init_threads() : c_int;
    return fftw_init_threads();
  }
  proc cleanup_threads() {
    extern proc fftw_cleanup_threads();
    fftw_cleanup_threads();
  }
  proc plan_with_nthreads(n : c_int) {
    extern proc fftw_plan_with_nthreads(n : c_int);
    fftw_plan_with_nthreads(n);
  }

  // Utilities -- not in FFTW
  proc abs((r,c) : fftw_complex) : real(64) {
    return sqrt(r*r + c*c);
  }

  proc re((r,c) : fftw_complex) : real(64) {
    return r;
  }
  proc im((r,c) : fftw_complex) : real(64) {
    return c;
  }

}


Oops, this should have gone in with the last commit

git-svn-id: 88467cb1fb04b8a755be7e1ee1026be4190196ef@18380 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

use BlockDist, CyclicDist, BlockCycDist;

enum DistType { default, block, cyclic, blockcyclic };

config param distType: DistType = DistType.default;

config const n1 = 100;
config const n2 = 20;
config const n3 = 5;
config const n4 = 3;
config const n5 = max(int(64));
config const o5 = 8;

const Space1 = [1..n1];
const Space2 = [1..n2, 1..n2];
const Space3 = [1..n3, 1..n3, 1..n3];
const Space4 = [1..n4, 1..n4, 1..n4, 1..n4];
const Space2D64 = [n5-o5..n5, n5-o5..n5];

proc setupDistributions() {
  if distType == DistType.default {
    return (
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist())
           );
  }
  if distType == DistType.block {
    return (
            new dmap(new Block(rank=1, boundingBox=Space1)),
            new dmap(new Block(rank=2, boundingBox=Space2)),
            new dmap(new Block(rank=3, boundingBox=Space3)),
            new dmap(new Block(rank=4, boundingBox=Space4)),
            new dmap(new Block(rank=2, idxType=int(64), boundingBox=Space2D64))
           );
  }
  if distType == DistType.cyclic {
    return (
            new dmap(new Cyclic(startIdx=0)),
            new dmap(new Cyclic(startIdx=(0,0))),
            new dmap(new Cyclic(startIdx=(0,0,0))),
            new dmap(new Cyclic(startIdx=(0,0,0,0))),
            new dmap(new Cyclic(startIdx=(0:int(64), 0:int(64))))
           );
  }
  if distType == DistType.blockcyclic {
    return (
            new dmap(new BlockCyclic(startIdx=tuple(0), blocksize=tuple(3))),
            new dmap(new BlockCyclic(startIdx=(0,0), blocksize=(3,3))),
            new dmap(new BlockCyclic(startIdx=(0,0,0), blocksize=(3,3,3))),
            new dmap(new BlockCyclic(startIdx=(0,0,0,0), blocksize=(3,3,3,3))),
            new dmap(new BlockCyclic(startIdx=(0:int(64),0:int(64)), blocksize=(2,3)))
           );
  }
}

const (Dist1D, Dist2D, Dist3D, Dist4D, Dist2D64) = setupDistributions();

//
// creates a tuple of size 'rank' initialized with values 'x'
//
proc fill(param rank, x) {
  if rank == 1 {
    var y: 1*x.type;
    y(1) = x;
    return y;
  } else if rank == 2 {
    return (x, x);
  } else {
    return (x, (...fill(rank-1, x)));
  }
}

//
// creates a domain of rank 'rank' with each dimension equal to 'extent'
//
proc rankDomain(param rank, extent) {
  return [(...fill(rank, extent))];
}

//
// returns a next integer starting at 1
//
var next_i = 0;
proc next() { next_i += 1; return next_i; }

use BlockDist, CyclicDist, BlockCycDist;

enum DistType { default, block, cyclic, blockcyclic };

config param distType: DistType = DistType.default;

config const n1 = 100;
config const n2 = 20;
config const n3 = 5;
config const n4 = 3;
config const n5 = max(int(64));
config const o5 = 8;

const Space1 = [1..n1];
const Space2 = [1..n2, 1..n2];
const Space3 = [1..n3, 1..n3, 1..n3];
const Space4 = [1..n4, 1..n4, 1..n4, 1..n4];
const Space2D64 = [n5-o5..n5, n5-o5..n5];

proc setupDistributions() {
  if distType == DistType.default {
    return (
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist()),
            new dmap(new DefaultDist())
           );
  }
  if distType == DistType.block {
    return (
            new dmap(new Block(rank=1, boundingBox=Space1)),
            new dmap(new Block(rank=2, boundingBox=Space2)),
            new dmap(new Block(rank=3, boundingBox=Space3)),
            new dmap(new Block(rank=4, boundingBox=Space4)),
            new dmap(new Block(rank=2, idxType=int(64), boundingBox=Space2D64))
           );
  }
  if distType == DistType.cyclic {
    return (
            new dmap(new Cyclic(startIdx=0)),
            new dmap(new Cyclic(startIdx=(0,0))),
            new dmap(new Cyclic(startIdx=(0,0,0))),
            new dmap(new Cyclic(startIdx=(0,0,0,0))),
            new dmap(new Cyclic(startIdx=(0:int(64), 0:int(64))))
           );
  }
  if distType == DistType.blockcyclic {
    return (
            new dmap(new BlockCyclic(rank=1, idxType=int, low=tuple(0), blk=tuple(3))),
            new dmap(new BlockCyclic(rank=2, idxType=int, low=(0,0), blk=(3,3))),
            new dmap(new BlockCyclic(rank=3, idxType=int, low=(0,0,0), blk=(3,3,3))),
            new dmap(new BlockCyclic(rank=4, idxType=int, low=(0,0,0,0), blk=(3,3,3,3))),
            new dmap(new BlockCyclic(rank=2, idxType=int(64), low=(0:int(64),0:int(64)), blk=(2,3)))
           );
  }
}

const (Dist1D, Dist2D, Dist3D, Dist4D, Dist2D64) = setupDistributions();

//
// creates a tuple of size 'rank' initialized with values 'x'
//
proc fill(param rank, x) {
  if rank == 1 {
    var y: 1*x.type;
    y(1) = x;
    return y;
  } else if rank == 2 {
    return (x, x);
  } else {
    return (x, (...fill(rank-1, x)));
  }
}

//
// creates a domain of rank 'rank' with each dimension equal to 'extent'
//
proc rankDomain(param rank, extent) {
  return [(...fill(rank, extent))];
}

//
// returns a next integer starting at 1
//
var next_i = 0;
proc next() { next_i += 1; return next_i; }