uts_distributed now correctly shows tree size

/*
**  UTS -- The Unbalanced Tree Search
**
**  James Dinan <dinan@cray.com>
**  July, 2007
*/

use Math;
use Time;
use sha1_rng;
use DistributedWorkQueue;
use Queue;

/**** UTS TYPES ****/
enum NodeDistrib { Binomial, Geometric, Hybrid, Constant };
enum GeoDistrib  { GeoFixed, GeoLinear, GeoPoly, GeoCyclic };

/**** COMPILE PARAMS ****/
config param debug:    bool = false;
config param parallel: bool = true;
param uts_version = "2.1";

/**** UTS CONFIG VARS ****/
config const SEED:        uint = 16;    // Seed for the hash of the root node
config const MAX_CHILDREN: int = 10;    // Max. Children a node may have
config const B_0:          int = 4;     // Branching Factor
config const MAX_DEPTH:    int = 6;     // Depth limit
config const nonLeafProb: real = 0.10;  // Probability of a non-leaf (binomial)
config const nonLeafBF:    int = 4;     // Non-leaf branching factor (binomial)
config const shiftDepth:  real = 0.5;   // Depth at which hybrid trees go from Geo=>Bin

config const distrib: NodeDistrib = NodeDistrib.Geometric;
config const geoDist: GeoDistrib  = GeoDistrib.GeoFixed;
config const testMode             = false;

// Global thread counter
var thread_cnt: atomic int;
var padding1 : 7 * int;
var threads_spawned: atomic int;
var padding2 : 7 * int;

// Global stats, updated by a thread when it exits
// FIXME: These will become hotspots
var global_count: atomic int;
var padding3 : 7 * int;
var global_maxDepth: atomic int;
var padding4 : 7 * int;
var global_elemsProcessed : atomic int;

/**** UTS TreeNode Class ****/
class TreeNode {
  var depth:   int;
  var hash:    [1..20] uint(8);

  var nChildren: int = 0;

  // Generate this node's children
  proc genChildren(q : Queue(TreeNode)): int {
    select distrib {
      when NodeDistrib.Geometric do
        nChildren = numGeoChildren(geoDist);
      when NodeDistrib.Binomial do
        nChildren = numBinChildren();
      when NodeDistrib.Hybrid do
        nChildren = numHybridChildren();
      when NodeDistrib.Constant do
        nChildren = numConstChildren();
    }

    if debug then writeln("Constructing ", nChildren, " children");

    for i in 0..nChildren-1 {
      if debug then writeln("  + (", depth, ", ", i, ")");
      var t = new TreeNode(depth+1);
      rng_spawn(hash[1], t.hash[1], i:sha_int);
      q.enqueue(t);
    }

    return nChildren;
  }

  // Constant Distribution: Find the number of children
  //  Note: This is a _balanced_ tree.
  proc numConstChildren(): int {
    return if depth < MAX_DEPTH then B_0 else 0;
  }

  // Hybrid trees are geometric at the top of the tree and binomial
  // below a certain depth.
  proc numHybridChildren(): int {
    if (depth < shiftDepth * MAX_DEPTH) then
      return numGeoChildren(geoDist);

    else
      return numBinChildren();
  }

  // Binomial Distribution: Find the number of children
  //  Note: distribution is identical everywhere below root
  proc numBinChildren(): int {
    if (depth == 0) then
      return B_0;

    else
      return if to_prob(rng_rand(hash[1])) < nonLeafProb then nonLeafBF else 0;
  }

  // Geometric Distribution: Find the number of children
  proc numGeoChildren(shape: GeoDistrib): int {
    var b_i: real = B_0;

    // use shape function to compute target b_i
    if (depth > 0) {
      select shape {

        // Expected size is polynomial in depth
        when GeoDistrib.GeoPoly do
          b_i = B_0 * (depth ** (-log(B_0:real)/log(MAX_DEPTH:real)));

        // Expected size is cyclic
        when GeoDistrib.GeoCyclic {
          if (depth > 5 * MAX_DEPTH) then
            b_i = 0.0;
          else
            b_i = B_0 ** sin(2.0*3.141592653589793*depth/MAX_DEPTH);
        }

        // Expected size is the same at all nodes up to max depth
        when GeoDistrib.GeoFixed do
          b_i = if (depth < MAX_DEPTH) then B_0 else 0;

        // Expected size decreases linearly in b_i
        when GeoDistrib.GeoLinear do
          b_i =  B_0 * (1.0 - depth:real / MAX_DEPTH:real);

      }
    }

    // given target b_i, find prob p so expected value of
    // geometric distribution is b_i.
    var p: real = 1.0 / (1.0 + b_i);

    // get uniform random number on [0,1)
    var u: real = to_prob(rng_rand(hash[1]));

    // max number of children at this cumulative probability
    // (from inverse geometric cumulative density function)
    return floor(log(1 - u) / log(1 - p)):int;
  }
}

/* Compute atomically:  this = max(this, other) */
proc atomic_int64.max(other: int) {
  var curMax = this.read();
  while curMax < other && !this.compareExchange(curMax, other) do
    curMax = this.read();
}

proc ratomic_int64.max(other: int) {
  on this {
  	var curMax = this.read();
  	while curMax < other && !this.compareExchange(curMax, other) do
    	curMax = this.read();
  }
}

/*
** Print out search parameters
*/
proc uts_showSearchParams() {
  writeln("UTS v", uts_version, " - Unbalanced Tree Search (",
      if parallel then "Parallel Chapel (DeQueue)" else "Sequential Chapel (DeQueue)", ")");

  writeln("Tree type: ", distrib);
  writeln("Tree shape parameters:");
  writeln("  root branching factor b_0 = ", B_0, ", root seed = ", SEED);

  if (distrib == NodeDistrib.Geometric || distrib == NodeDistrib.Hybrid) then
    writeln("  Geo. params: Max_Depth = ", MAX_DEPTH, ", shape function = ", geoDist);

  if (distrib == NodeDistrib.Binomial || distrib == NodeDistrib.Hybrid) then
    writeln("  Bin. params: nonLeafProb=", nonLeafProb, " nonLeafBF=", nonLeafBF);

  if (distrib == NodeDistrib.Hybrid) then
    writeln("  Hybrid params: shiftDepth=", shiftDepth);

  writeln("Random number generator: ", rng_getName());

  writeln();
}

proc main() {
  var t_create: Timer;
  var root: TreeNode;
  var queue = new DistributedWorkQueue(TreeNode);

  // Create the root and push it into a queue
  root = new TreeNode(0);
  rng_init(root.hash[1], SEED:sha_int);
  global_count.add(1);
  queue.enqueue(root);

  /*uts_showSearchParams();*/

  writeln("Performing ", if parallel then "parallel" else "serial", " tree search...");

  t_create.start();
  coforall loc in Locales do on loc {
    var localQueue = queue.localQueue;
    coforall tid in 0 .. #here.maxTaskPar {
      var (hasElem, elem) : (bool, TreeNode);
      var maxDepth : int;
      while global_count.read() != 0 {
        (hasElem, elem) = localQueue.dequeue();
        if !hasElem {
          chpl_task_yield();
          continue;
        }

        global_count.add(elem.genChildren(localQueue));
        global_elemsProcessed.add(1);
        global_count.sub(1);

        if elem.depth > maxDepth {
          global_maxDepth.max(elem.depth);
          maxDepth = global_maxDepth.read();
        }

        delete elem;
      }
    }
  }
  t_create.stop();

  writeln();
  if !testMode then writeln("Threads spawned: ", threads_spawned.read());
  writeln("Tree size = ", global_elemsProcessed.read(),
            ", depth = ", global_maxDepth.read());
  if !testMode then writeln("Time: t_create= ", t_create.elapsed(),
          " (", global_elemsProcessed.read()/t_create.elapsed(), " nodes/sec)");

}

/*
**  UTS -- The Unbalanced Tree Search
**
**  James Dinan <dinan@cray.com>
**  July, 2007
*/

use Math;
use Time;
use sha1_rng;
use DistributedWorkQueue;
use Queue;

/**** UTS TYPES ****/
enum NodeDistrib { Binomial, Geometric, Hybrid, Constant };
enum GeoDistrib  { GeoFixed, GeoLinear, GeoPoly, GeoCyclic };

/**** COMPILE PARAMS ****/
config param debug:    bool = false;
config param parallel: bool = true;
param uts_version = "2.1";

/**** UTS CONFIG VARS ****/
config const SEED:        uint = 16;    // Seed for the hash of the root node
config const MAX_CHILDREN: int = 10;    // Max. Children a node may have
config const B_0:          int = 4;     // Branching Factor
config const MAX_DEPTH:    int = 6;     // Depth limit
config const nonLeafProb: real = 0.10;  // Probability of a non-leaf (binomial)
config const nonLeafBF:    int = 4;     // Non-leaf branching factor (binomial)
config const shiftDepth:  real = 0.5;   // Depth at which hybrid trees go from Geo=>Bin

config const distrib: NodeDistrib = NodeDistrib.Geometric;
config const geoDist: GeoDistrib  = GeoDistrib.GeoFixed;
config const testMode             = false;

// Global thread counter
var thread_cnt: atomic int;
var padding1 : 7 * int;
var threads_spawned: atomic int;
var padding2 : 7 * int;

// Global stats, updated by a thread when it exits
// FIXME: These will become hotspots
var global_count: atomic int;
var padding3 : 7 * int;
var global_maxDepth: atomic int;
var padding4 : 7 * int;
var global_elemsProcessed : atomic int;

/**** UTS TreeNode Class ****/
class TreeNode {
  var depth:   int;
  var hash:    [1..20] uint(8);

  var nChildren: int = 0;

  // Generate this node's children
  proc genChildren(q : Queue(TreeNode)): int {
    select distrib {
      when NodeDistrib.Geometric do
        nChildren = numGeoChildren(geoDist);
      when NodeDistrib.Binomial do
        nChildren = numBinChildren();
      when NodeDistrib.Hybrid do
        nChildren = numHybridChildren();
      when NodeDistrib.Constant do
        nChildren = numConstChildren();
    }

    if debug then writeln("Constructing ", nChildren, " children");

    for i in 0..nChildren-1 {
      if debug then writeln("  + (", depth, ", ", i, ")");
      var t = new TreeNode(depth+1);
      rng_spawn(hash[1], t.hash[1], i:sha_int);
      q.enqueue(t);
    }

    return nChildren;
  }

  // Constant Distribution: Find the number of children
  //  Note: This is a _balanced_ tree.
  proc numConstChildren(): int {
    return if depth < MAX_DEPTH then B_0 else 0;
  }

  // Hybrid trees are geometric at the top of the tree and binomial
  // below a certain depth.
  proc numHybridChildren(): int {
    if (depth < shiftDepth * MAX_DEPTH) then
      return numGeoChildren(geoDist);

    else
      return numBinChildren();
  }

  // Binomial Distribution: Find the number of children
  //  Note: distribution is identical everywhere below root
  proc numBinChildren(): int {
    if (depth == 0) then
      return B_0;

    else
      return if to_prob(rng_rand(hash[1])) < nonLeafProb then nonLeafBF else 0;
  }

  // Geometric Distribution: Find the number of children
  proc numGeoChildren(shape: GeoDistrib): int {
    var b_i: real = B_0;

    // use shape function to compute target b_i
    if (depth > 0) {
      select shape {

        // Expected size is polynomial in depth
        when GeoDistrib.GeoPoly do
          b_i = B_0 * (depth ** (-log(B_0:real)/log(MAX_DEPTH:real)));

        // Expected size is cyclic
        when GeoDistrib.GeoCyclic {
          if (depth > 5 * MAX_DEPTH) then
            b_i = 0.0;
          else
            b_i = B_0 ** sin(2.0*3.141592653589793*depth/MAX_DEPTH);
        }

        // Expected size is the same at all nodes up to max depth
        when GeoDistrib.GeoFixed do
          b_i = if (depth < MAX_DEPTH) then B_0 else 0;

        // Expected size decreases linearly in b_i
        when GeoDistrib.GeoLinear do
          b_i =  B_0 * (1.0 - depth:real / MAX_DEPTH:real);

      }
    }

    // given target b_i, find prob p so expected value of
    // geometric distribution is b_i.
    var p: real = 1.0 / (1.0 + b_i);

    // get uniform random number on [0,1)
    var u: real = to_prob(rng_rand(hash[1]));

    // max number of children at this cumulative probability
    // (from inverse geometric cumulative density function)
    return floor(log(1 - u) / log(1 - p)):int;
  }
}

/* Compute atomically:  this = max(this, other) */
proc atomic_int64.max(other: int) {
  var curMax = this.read();
  while curMax < other && !this.compareExchange(curMax, other) do
    curMax = this.read();
}

proc ratomic_int64.max(other: int) {
  on this {
  	var curMax = this.read();
  	while curMax < other && !this.compareExchange(curMax, other) do
    	curMax = this.read();
  }
}

/*
** Print out search parameters
*/
proc uts_showSearchParams() {
  writeln("UTS v", uts_version, " - Unbalanced Tree Search (",
      if parallel then "Parallel Chapel (DeQueue)" else "Sequential Chapel (DeQueue)", ")");

  writeln("Tree type: ", distrib);
  writeln("Tree shape parameters:");
  writeln("  root branching factor b_0 = ", B_0, ", root seed = ", SEED);

  if (distrib == NodeDistrib.Geometric || distrib == NodeDistrib.Hybrid) then
    writeln("  Geo. params: Max_Depth = ", MAX_DEPTH, ", shape function = ", geoDist);

  if (distrib == NodeDistrib.Binomial || distrib == NodeDistrib.Hybrid) then
    writeln("  Bin. params: nonLeafProb=", nonLeafProb, " nonLeafBF=", nonLeafBF);

  if (distrib == NodeDistrib.Hybrid) then
    writeln("  Hybrid params: shiftDepth=", shiftDepth);

  writeln("Random number generator: ", rng_getName());

  writeln();
}

proc main() {
  var t_create: Timer;
  var root: TreeNode;
  var queue = new DistributedWorkQueue(TreeNode);

  // Create the root and push it into a queue
  root = new TreeNode(0);
  rng_init(root.hash[1], SEED:sha_int);
  global_count.add(1);
  queue.enqueue(root);

  /*uts_showSearchParams();*/

  writeln("Performing ", if parallel then "parallel" else "serial", " tree search...");

  t_create.start();
  coforall loc in Locales do on loc {
    var localQueue = queue.localQueue;
    coforall tid in 0 .. #here.maxTaskPar {
      var (hasElem, elem) : (bool, TreeNode);
      var maxDepth : int;
      while global_count.read() != 0 {
        (hasElem, elem) = localQueue.dequeue();
        if !hasElem {
          chpl_task_yield();
          continue;
        }

        global_count.add(elem.genChildren(localQueue));
        global_elemsProcessed.add(1);
        global_count.sub(1);

        if elem.depth > maxDepth {
          global_maxDepth.max(elem.depth);
          maxDepth = global_maxDepth.read();
        }

        delete elem;
      }
    }
  }
  t_create.stop();

  writeln();
  if !testMode then writeln("Threads spawned: ", threads_spawned.read());
  writeln("Tree size = ", global_count.read(),
            ", depth = ", global_maxDepth.read());
  if !testMode then writeln("Time: t_create= ", t_create.elapsed(),
          " (", global_elemsProcessed.read()/t_create.elapsed(), " nodes/sec)");

}


Broke declaration of distribution and domain in stream back into two lines.
Combining this declaration into one line destroys performance (4.3 GB/s -> 0.3
GB/s).  I'm guessing this is because the finicky 'aligned' optimization does
not fire in that case.

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

//
// Use standard modules for Block distributions, Timing routines, Type
// utility functions, and Random numbers
//
use BlockDist, Time, Types, Random;

//
// Use shared user module for computing HPCC problem sizes
//
use HPCCProblemSize;

//
// The number of vectors and element type of those vectors
//
const numVectors = 3;
type elemType = real(64);

//
// Configuration constants to set the problem size (m) and the scalar
// multiplier, alpha
//
config const m = computeProblemSize(numVectors, elemType),
             alpha = 3.0;

//
// Configuration constants to set the number of trials to run and the
// amount of error to permit in the verification
//
config const numTrials = 10,
             epsilon = 0.0;

//
// Configuration constants to indicate whether or not to use a
// pseudo-random seed (based on the clock) or a fixed seed; or to
// specify the seed explicitly
//
config const useRandomSeed = true,
             seed = if useRandomSeed then SeedGenerator.clockMS else 314159265;

//
// Configuration constants to control what's printed -- benchmark
// parameters, input and output arrays, and/or statistics
//
config const printParams = true,
             printArrays = false,
             printStats = true;

//
// The program entry point
//
def main() {
  printConfiguration();   // print the problem size, number of trials, etc.

  //
  // ProblemSpace describes the index set for the three vectors.  It
  // is a 1D domain indexed using 64-bit ints that is distributed by
  // blocking the indices 1..m between the Locales, and it contains the
  // indices 1..m.
  //
  const BlockDist = new Block1D(bbox=[1..m]);
  const ProblemSpace: domain(1, int(64)) distributed BlockDist = [1..m];

  //
  // A, B, and C are the three distributed vectors, declared to store
  // a variable of type elemType for each index in ProblemSpace.
  //
  var A, B, C: [ProblemSpace] elemType;

  initVectors(B, C);    // Initialize the input vectors, B and C

  var execTime: [1..numTrials] real;                 // an array of timings

  for trial in 1..numTrials {                        // loop over the trials
    const startTime = getCurrentTime();              // capture the start time

    //
    // The main loop: Iterate over the vectors A, B, and C in a
    // parallel, zippered manner storing the elements as a, b, and c.
    // Compute the multiply-add on b and c, storing the result to a.
    //
    forall (a, b, c) in (A, B, C) do
      a = b + alpha * c;

    execTime(trial) = getCurrentTime() - startTime;  // store the elapsed time
  }

  const validAnswer = verifyResults(A, B, C);        // verify...
  printResults(validAnswer, execTime);               // ...and print the results
}

//
// Print the problem size and number of trials
//
def printConfiguration() {
  if (printParams) {
    printProblemSize(elemType, numVectors, m);
    writeln("Number of trials = ", numTrials, "\n");
  }
}

//
// Initialize vectors B and C using a random stream of values and
// optionally print them to the console
//
def initVectors(B, C) {
  var randlist = new RandomStream(seed);

  randlist.fillRandom(B);
  randlist.fillRandom(C);

  if (printArrays) {
    writeln("B is: ", B, "\n");
    writeln("C is: ", C, "\n");
  }
}

//
// Verify that the computation is correct
//
def verifyResults(A, B, C) {
  if (printArrays) then writeln("A is: ", A, "\n");  // optionally print A

  //
  // Compute the infinity-norm by computing the maximum reduction of the
  // absolute value of A's elements minus the multiply-add of B and C.
  // "[i in I]" represents an expression-level loop: "forall i in I"
  //
  const infNorm = max reduce [(a,b,c) in (A,B,C)] abs(a - (b + alpha * c));

  return (infNorm <= epsilon);    // return whether the error is acceptable
}

//
// Print out success/failure, the timings, and the GB/s value
//
def printResults(successful, execTimes) {
  writeln("Validation: ", if successful then "SUCCESS" else "FAILURE");
  if (printStats) {
    const totalTime = + reduce execTimes,
          avgTime = totalTime / numTrials,
          minTime = min reduce execTimes;
    writeln("Execution time:");
    writeln("  tot = ", totalTime);
    writeln("  avg = ", avgTime);
    writeln("  min = ", minTime);

    const GBPerSec = numVectors * numBytes(elemType) * (m / minTime) * 1e-9;
    writeln("Performance (GB/s) = ", GBPerSec);
  }
}

//
// Use standard modules for Block distributions, Timing routines, Type
// utility functions, and Random numbers
//
use BlockDist, Time, Types, Random;

//
// Use shared user module for computing HPCC problem sizes
//
use HPCCProblemSize;

//
// The number of vectors and element type of those vectors
//
const numVectors = 3;
type elemType = real(64);

//
// Configuration constants to set the problem size (m) and the scalar
// multiplier, alpha
//
config const m = computeProblemSize(numVectors, elemType),
             alpha = 3.0;

//
// Configuration constants to set the number of trials to run and the
// amount of error to permit in the verification
//
config const numTrials = 10,
             epsilon = 0.0;

//
// Configuration constants to indicate whether or not to use a
// pseudo-random seed (based on the clock) or a fixed seed; or to
// specify the seed explicitly
//
config const useRandomSeed = true,
             seed = if useRandomSeed then SeedGenerator.clockMS else 314159265;

//
// Configuration constants to control what's printed -- benchmark
// parameters, input and output arrays, and/or statistics
//
config const printParams = true,
             printArrays = false,
             printStats = true;

//
// The program entry point
//
def main() {
  printConfiguration();   // print the problem size, number of trials, etc.

  //
  // ProblemSpace describes the index set for the three vectors.  It
  // is a 1D domain indexed using 64-bit ints that is distributed by
  // blocking the indices 1..m between the Locales, and it contains the
  // indices 1..m.
  //
  const ProblemSpace: domain(1, int(64)) distributed new Block1D(bbox=[1..m])
                    = [1..m];

  //
  // A, B, and C are the three distributed vectors, declared to store
  // a variable of type elemType for each index in ProblemSpace.
  //
  var A, B, C: [ProblemSpace] elemType;

  initVectors(B, C);    // Initialize the input vectors, B and C

  var execTime: [1..numTrials] real;                 // an array of timings

  for trial in 1..numTrials {                        // loop over the trials
    const startTime = getCurrentTime();              // capture the start time

    //
    // The main loop: Iterate over the vectors A, B, and C in a
    // parallel, zippered manner storing the elements as a, b, and c.
    // Compute the multiply-add on b and c, storing the result to a.
    //
    forall (a, b, c) in (A, B, C) do
      a = b + alpha * c;

    execTime(trial) = getCurrentTime() - startTime;  // store the elapsed time
  }

  const validAnswer = verifyResults(A, B, C);        // verify...
  printResults(validAnswer, execTime);               // ...and print the results
}

//
// Print the problem size and number of trials
//
def printConfiguration() {
  if (printParams) {
    printProblemSize(elemType, numVectors, m);
    writeln("Number of trials = ", numTrials, "\n");
  }
}

//
// Initialize vectors B and C using a random stream of values and
// optionally print them to the console
//
def initVectors(B, C) {
  var randlist = new RandomStream(seed);

  randlist.fillRandom(B);
  randlist.fillRandom(C);

  if (printArrays) {
    writeln("B is: ", B, "\n");
    writeln("C is: ", C, "\n");
  }
}

//
// Verify that the computation is correct
//
def verifyResults(A, B, C) {
  if (printArrays) then writeln("A is: ", A, "\n");  // optionally print A

  //
  // Compute the infinity-norm by computing the maximum reduction of the
  // absolute value of A's elements minus the multiply-add of B and C.
  // "[i in I]" represents an expression-level loop: "forall i in I"
  //
  const infNorm = max reduce [(a,b,c) in (A,B,C)] abs(a - (b + alpha * c));

  return (infNorm <= epsilon);    // return whether the error is acceptable
}

//
// Print out success/failure, the timings, and the GB/s value
//
def printResults(successful, execTimes) {
  writeln("Validation: ", if successful then "SUCCESS" else "FAILURE");
  if (printStats) {
    const totalTime = + reduce execTimes,
          avgTime = totalTime / numTrials,
          minTime = min reduce execTimes;
    writeln("Execution time:");
    writeln("  tot = ", totalTime);
    writeln("  avg = ", avgTime);
    writeln("  min = ", minTime);

    const GBPerSec = numVectors * numBytes(elemType) * (m / minTime) * 1e-9;
    writeln("Performance (GB/s) = ", GBPerSec);
  }
}


Code to do full paircounts

Code to do full paircounts

use Histogram;
use Time;

// Test flags
config const isTest=false;

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

//Tree parameters
param NDIM  = 3;
param NEXT=2; // w, r2
param NTOT=NDIM+NEXT;
param W=NDIM;
param R2=NDIM+1;
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() {
  if isTest {
    testPairs();
  } else {
    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, _n2,_ndx : [Dpart] int;

  proc Particle3D(npart1 : int) {
    npart = npart1;
    Darr = {0.. #NTOT, 0.. #npart};
    Dpart = {0.. #npart};
  }

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

}

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;
  }

  return pp;
}

proc splitOn(pp : Particle3D, idom : range, 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;
    }
  }

  for kk in 0.. #NTOT {
    forall ii in idom {
      pp._tmp[pp._ndx[ii]] = pp.arr[kk,ii];
    }
    forall ii in idom {
      pp.arr[kk,ii] = pp._tmp[ii];
    }
  }

  return lnpart;
}

class KDNode {
  var lo, hi,npart,id : int;
  var dom : domain(1);
  var xcen : [0.. #NDIM]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 : [0.. #NDIM] real;
  for idim in 0.. #NDIM {
    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 0.. #NDIM {
      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 : [0.. #NDIM] real;
  dx = pmax - pmin;
  var splitDim = 0;
  for idim in 0.. #NDIM {
    if (dx[idim] > dx[splitDim]) then splitDim=idim;
  }

  // Split
  var lnpart = splitOn(pp, lo..hi,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;
  }

}

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 testPairs() {
  var pp = readFile("test.dat");

  // 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
  smuAccumulate(hh,pp,pp,pp.Dpart,pp.Dpart,1.0);
  for xx in hh.bins(1) do writef("%12.4dr",xx);
  writeln();
  for xx in hh.bins(2) do writef("%12.4dr",xx);
  writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      writef("%20.5er ",hh[(ii,jj)]);
    }
    writeln();
  }

  // Build the tree
  var root1 = BuildTree(pp, 0, pp.npart-1, 0);
  hh.reset();
  sync TreeAccumulate(hh,pp,pp,root1,root1);
  for xx in hh.bins(1) do writef("%12.4dr",xx);
  writeln();
  for xx in hh.bins(2) do writef("%12.4dr",xx);
  writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      writef("%20.5er ",hh[(ii,jj)]);
    }
    writeln();
  }

}

proc doPairs() {
  var tt : Timer;

  // Read in the file
  tt.clear(); tt.start();
  var pp1 = readFile(fn1);
  var pp2 = readFile(fn2);
  tt.stop();
  writef("Read in %i lines from file %s \n", pp1.npart, fn1);
  writef("Read in %i lines from file %s \n", pp2.npart, fn2);
  writef("Time to read : %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();
  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();
  writef("Time to tree paircount : %r \n", tt.elapsed());
  //writeHist("%s.tree".format(pairfn),hh);
}

use Histogram;
use Time;

// Test flags
config const isTest=false;

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

//Tree parameters
param NDIM  = 3;
param NEXT=2; // w, r2
param NTOT=NDIM+NEXT;
param W=NDIM;
param R2=NDIM+1;
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() {
  testPairs();
}

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

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

  proc Particle3D(npart1 : int) {
    npart = npart1;
    Darr = {0.. #NTOT, 0.. #npart};
    Dpart = {0.. #npart};
  }

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

}

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;
  }

  return pp;
}

proc splitOn(pp : Particle3D, idom : range, 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;
    }
  }

  for kk in 0.. #NTOT {
    forall ii in idom {
      pp._tmp[pp._ndx[ii]] = pp.arr[kk,ii];
    }
    forall ii in idom {
      pp.arr[kk,ii] = pp._tmp[ii];
    }
  }

  return lnpart;
}

class KDNode {
  var lo, hi,npart,id : int;
  var dom : domain(1);
  var xcen : [0.. #NDIM]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 : [0.. #NDIM] real;
  for idim in 0.. #NDIM {
    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 0.. #NDIM {
      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 : [0.. #NDIM] real;
  dx = pmax - pmin;
  var splitDim = 0;
  for idim in 0.. #NDIM {
    if (dx[idim] > dx[splitDim]) then splitDim=idim;
  }

  // Split
  var lnpart = splitOn(pp, lo..hi,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;
  }

}

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 testPairs() {
  var pp = readFile("test.dat");

  // 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
  smuAccumulate(hh,pp,pp,pp.Dpart,pp.Dpart,1.0);
  for xx in hh.bins(1) do writef("%12.4dr",xx);
  writeln();
  for xx in hh.bins(2) do writef("%12.4dr",xx);
  writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      writef("%20.5er ",hh[(ii,jj)]);
    }
    writeln();
  }

  // Build the tree
  var root1 = BuildTree(pp, 0, pp.npart-1, 0);
  hh.reset();
  sync TreeAccumulate(hh,pp,pp,root1,root1);
  for xx in hh.bins(1) do writef("%12.4dr",xx);
  writeln();
  for xx in hh.bins(2) do writef("%12.4dr",xx);
  writeln("\n##");
  for ii in hh.Dhist.dim(1) {
    for jj in hh.Dhist.dim(2) {
      writef("%20.5er ",hh[(ii,jj)]);
    }
    writeln();
  }

}