Initialize root locale in DefaultArchitecture
to correct for some testing instability.

git-svn-id: 5aabc979e906cfa1cd4a03d781c193b99fe6e46f@21415 3a8e244f-b0f2-452b-bcba-4c88e055c3ca

// DefaultArchitecture.chpl
//
// Provides a default architectural description.
// This architectural description is backward compatible with
// the architecture implicitly provided by releases 1.6 and preceding.
//
pragma "no use ChapelStandard"
module DefaultArchitecture {

  use ChapelLocale;
  use DefaultRectangular;
  use ChapelNumLocales;
  use RootLocale;
  use Sys;

  // We would really like a class-static storage class.(C++ nomenclature)
  var doneCreatingLocales: bool = false;

  //
  // A concrete class representing the nodes in this architecture.
  //
  class DefaultNode : locale {
    const callStackSize: int;
    const _node_id : int;
    const local_name : string;

    // This constructor must be invoked "on" the node
    // that it is intended to represent.  This trick is used
    // to establish the equivalence the "locale" field of the locale object
    // and the node ID portion of any wide pointer referring to it.
    proc DefaultNode() {
      if doneCreatingLocales {
        halt("Cannot create additional DefaultNode instances");
      }
      init;
    }

    proc DefaultNode(parent_loc : locale) {
      parent = parent_loc;
      init;
    }

    // chpl_nodeID is the node ID associated with the running image.
    proc chpl_id() return _node_id;
    proc chpl_name() return local_name;

    proc readWriteThis(f) {
      // Most classes will define it like this:
//      f <~> name;
      // but here it is defined thus for backward compatibility.
      f <~> new ioLiteral("LOCALE") <~> _node_id;
    }

    proc getChildSpace() return emptyLocaleSpace;

    proc getChildCount() return 0;

    iter getChildIndices() : int {
      for idx in emptyLocaleSpace do
        yield idx;
    }

    proc getChild(idx:int) : locale {
      // This is a temporary implementation, where an index of zero is forced
      // to mean "here".  A better solution is on the way.
      // TODO: Run index lookup through rootLocale.findLocale(lid:chpl_localeID_t);
      if idx == 0 then return this;
      else
        if boundsChecking then
          halt("requesting a child from a DefaultNode locale");
      return nil;
    }

    iter getChldren() : locale  {
      for loc in emptyLocales do
        yield loc;
    }

    proc getChildArray() {
      return emptyLocales;
    }

    // Part of the public interface required by the compiler
    // These are dynamically dispatched, so cannot be inlined.
    proc taskInit() {}
    proc taskExit() {}

    //------------------------------------------------------------------------{
    //- Implementation (private)
    //-
    proc init {
      _node_id = __primitive("cast", int(64), __primitive("chpl_nodeID"));

      // chpl_nodeName is defined in chplsys.c.
      // It supplies a node name obtained by running uname(3) on the current node.
      // For this reason (as well), the constructor (or at least this init method)
      // must be run on the node it is intended to describe.
//      extern proc getenv(s:string) : string;
      var comm, spawnfn : string;
      extern proc chpl_nodeName() : string;
      // sys_getenv returns zero on success.
      if sys_getenv("CHPL_COMM", comm) == 0 && comm == "gasnet" &&
        sys_getenv("GASNET_SPAWNFN", spawnfn) == 0 && spawnfn == "L"
      then local_name = chpl_nodeName() + "-" + _node_id : string;
      else local_name = chpl_nodeName();

      extern proc chpl_task_getCallStackSize(): int;
      callStackSize = chpl_task_getCallStackSize();

      extern proc chpl_numCoresOnThisLocale(): int;
      numCores = chpl_numCoresOnThisLocale();
    }
    //------------------------------------------------------------------------}
  }

  //
  // An instance of this class is the default contents 'rootLocale'.
  //
  // In the current implementation a platform-specific architectural description
  // may overwrite this instance or any of its children to establish a more customized
  // representation of the system resources.
  //
  class DefaultRootLocale : RootLocale {

    // Would like to make myLocaleSpace distributed with one index per node.
    const myLocaleSpace: domain(1) = {0..numLocales-1};
    const myLocales: [myLocaleSpace] locale;

    proc DefaultRootLocale()
    {
      // A bootstrap routine that returns (chpl_localeID_t){.node = <arg>, .subloc = 0}.
      extern proc chpl_return_localeID_node(node:int(32)) : chpl_localeID_t;

      parent = nil;
      numCores = 0;

      // We cannot use a forall here because the default leader iterator will
      // access 'Locales' and 'here', which are not yet initialized.
      for locIdx in myLocaleSpace
      {
        var locID = chpl_return_localeID_node(locIdx:int(32));
        on __primitive("chpl_on_locale_num", locID)
        {
          // chpl_on_locale_num sets the localeID portion of "here", but
          // leaves the addr portion as a NULL pointer.
          const node = new DefaultNode(this);

          // So immediately after creating a new locale to represent here, we
          // have to call this primitive to insert it into task-private storage.
          __primitive("_task_set_here", node);

          myLocales[locIdx] = node;
          numCores += node.numCores;
        }
      }
      doneCreatingLocales = true;

      // Sanity check.
      // Ensure that the locale representing node x has a nodeID of x.
      // At the top level, this lets us look up the locale representing the node on which
      // an object o lives by executing
      //  myLocales[__primitive("_wide_get_node", o)].
      for locIdx in myLocaleSpace {
        var loc = myLocales[locIdx];
        if __primitive("_wide_get_node", loc) != locIdx then
          halt("In this architecture, we expect the locale whose index is x to live on node x.");
      }

      // Programs traditionally expect the startup process to run on Locales[0], so
      // this is how we mimic that behavior.
      // Note that this means we have to ask for here.parent or rootLocale to get
      // the root locale of the default architecture.
      var loc = myLocales[0];
      __primitive("_task_set_here", loc);
      var locID = chpl_return_localeID_node(0:int(32));
      __primitive("_task_set_locale", locID);
    }

    // Has to be globally unique and not equal to a node ID.
    // We return numLocales for now, since we expect nodes to be numbered less than this.
    // -1 is used in the abstract locale class to specify an invalid node ID.
    proc chpl_id() return numLocales;
    proc chpl_name() return local_name();
    proc local_name() return "rootLocale";

    proc readWriteThis(f) {
      f <~> name;
    }

    proc getChildCount() return this.myLocaleSpace.numIndices;

    proc getChildSpace() return this.myLocaleSpace;

    iter getChildIndices() : int {
      for idx in this.myLocaleSpace do
        yield idx;
    }

    proc getChild(idx:int) return this.myLocales[idx];

    iter getChldren() : locale  {
      for loc in this.myLocales do
        yield loc;
    }

    proc getDefaultLocaleSpace() return this.myLocaleSpace;
    proc getDefaultLocaleArray() return myLocales;

    proc localeIDtoLocale(id : chpl_localeID_t) {
      // In the default architecture, there are only nodes and no sublocales.
      // What is more, the nodeID portion of a wide pointer is the same as
      // the index into myLocales that yields the locale representing that node.
      return myLocales[id.node];
    }
  }

  if (rootLocale == nil) then
    rootLocale = new DefaultRootLocale();

  // Expose the underlying locales array (and its domain)
  // for user convenience and backward compatibility.
  // The downcase is because we cannot move the domain and array return types
  // into the base (abstract) locale class.
  // That would make locales depend on arrays which depend on locales....
  // If we had a way to express domain and array return types abstractly,
  // that problem would go away. <hilde>
  const Locales => (rootLocale:RootLocale).getDefaultLocaleArray();
  const LocaleSpace = Locales.domain;
}

// DefaultArchitecture.chpl
//
// Provides a default architectural description.
// This architectural description is backward compatible with
// the architecture implicitly provided by releases 1.6 and preceding.
//
pragma "no use ChapelStandard"
module DefaultArchitecture {

  use ChapelLocale;
  use DefaultRectangular;
  use ChapelNumLocales;
  use RootLocale;
  use Sys;

  // We would really like a class-static storage class.(C++ nomenclature)
  var doneCreatingLocales: bool = false;

  //
  // A concrete class representing the nodes in this architecture.
  //
  class DefaultNode : locale {
    const callStackSize: int;
    const _node_id : int;
    const local_name : string;

    // This constructor must be invoked "on" the node
    // that it is intended to represent.  This trick is used
    // to establish the equivalence the "locale" field of the locale object
    // and the node ID portion of any wide pointer referring to it.
    proc DefaultNode() {
      if doneCreatingLocales {
        halt("Cannot create additional DefaultNode instances");
      }
      init;
    }

    proc DefaultNode(parent_loc : locale) {
      parent = parent_loc;
      init;
    }

    // chpl_nodeID is the node ID associated with the running image.
    proc chpl_id() return _node_id;
    proc chpl_name() return local_name;

    proc readWriteThis(f) {
      // Most classes will define it like this:
//      f <~> name;
      // but here it is defined thus for backward compatibility.
      f <~> new ioLiteral("LOCALE") <~> _node_id;
    }

    proc getChildSpace() return emptyLocaleSpace;

    proc getChildCount() return 0;

    iter getChildIndices() : int {
      for idx in emptyLocaleSpace do
        yield idx;
    }

    proc getChild(idx:int) : locale {
      // This is a temporary implementation, where an index of zero is forced
      // to mean "here".  A better solution is on the way.
      // TODO: Run index lookup through rootLocale.findLocale(lid:chpl_localeID_t);
      if idx == 0 then return this;
      else
        if boundsChecking then
          halt("requesting a child from a DefaultNode locale");
      return nil;
    }

    iter getChldren() : locale  {
      for loc in emptyLocales do
        yield loc;
    }

    proc getChildArray() {
      return emptyLocales;
    }

    // Part of the public interface required by the compiler
    // These are dynamically dispatched, so cannot be inlined.
    proc taskInit() {}
    proc taskExit() {}

    //------------------------------------------------------------------------{
    //- Implementation (private)
    //-
    proc init {
      _node_id = __primitive("cast", int(64), __primitive("chpl_nodeID"));

      // chpl_nodeName is defined in chplsys.c.
      // It supplies a node name obtained by running uname(3) on the current node.
      // For this reason (as well), the constructor (or at least this init method)
      // must be run on the node it is intended to describe.
//      extern proc getenv(s:string) : string;
      var comm, spawnfn : string;
      extern proc chpl_nodeName() : string;
      // sys_getenv returns zero on success.
      if sys_getenv("CHPL_COMM", comm) == 0 && comm == "gasnet" &&
        sys_getenv("GASNET_SPAWNFN", spawnfn) == 0 && spawnfn == "L"
      then local_name = chpl_nodeName() + "-" + _node_id : string;
      else local_name = chpl_nodeName();

      extern proc chpl_task_getCallStackSize(): int;
      callStackSize = chpl_task_getCallStackSize();

      extern proc chpl_numCoresOnThisLocale(): int;
      numCores = chpl_numCoresOnThisLocale();
    }
    //------------------------------------------------------------------------}
  }

  //
  // An instance of this class is the default contents 'rootLocale'.
  //
  // In the current implementation a platform-specific architectural description
  // may overwrite this instance or any of its children to establish a more customized
  // representation of the system resources.
  //
  class DefaultRootLocale : RootLocale {

    // Would like to make myLocaleSpace distributed with one index per node.
    const myLocaleSpace: domain(1) = {0..numLocales-1};
    const myLocales: [myLocaleSpace] locale;

    proc DefaultRootLocale()
    {
      parent = nil;
      numCores = 0;

      // We cannot use a forall here because the default leader iterator will
      // access 'Locales' and 'here', which are not yet initialized.
      for locIdx in myLocaleSpace
      {
        // A bootstrap routine that returns (chpl_localeID_t){.node = <arg>, .subloc = 0}.
        extern proc chpl_return_localeID_node(node:int(32)) : chpl_localeID_t;

        var locID = chpl_return_localeID_node(locIdx:int(32));
        on __primitive("chpl_on_locale_num", locID)
        {
          // chpl_on_locale_num sets the localeID portion of "here", but
          // leaves the addr portion as a NULL pointer.
          const node = new DefaultNode(this);

          // So immediately after creating a new locale to represent here, we
          // have to call this primitive to insert it into task-private storage.
          __primitive("_task_set_here", node);

          myLocales[locIdx] = node;
          numCores += node.numCores;
        }
      }
      doneCreatingLocales = true;

      // Sanity check.
      // Ensure that the locale representing node x has a nodeID of x.
      // At the top level, this lets us look up the locale representing the node on which
      // an object o lives by executing
      //  myLocales[__primitive("_wide_get_node", o)].
      for locIdx in myLocaleSpace {
        var loc = myLocales[locIdx];
        if __primitive("_wide_get_node", loc) != locIdx then
          halt("In this architecture, we expect the locale whose index is x to live on node x.");
      }

      // Programs traditionally expect the startup process to run on Locales[0], so
      // this is how we mimic that behavior.
      // Note that this means we have to ask for here.parent or rootLocale to get
      // the root locale of the default architecture.
      var loc = myLocales[0];
      __primitive("_task_set_here", loc);
      var locID : chpl_localeID_t; // Init'd to {0,0}.
      __primitive("_task_set_locale", locID);
    }

    // Has to be globally unique and not equal to a node ID.
    // We return numLocales for now, since we expect nodes to be numbered less than this.
    // -1 is used in the abstract locale class to specify an invalid node ID.
    proc chpl_id() return numLocales;
    proc chpl_name() return local_name();
    proc local_name() return "rootLocale";

    proc readWriteThis(f) {
      f <~> name;
    }

    proc getChildCount() return this.myLocaleSpace.numIndices;

    proc getChildSpace() return this.myLocaleSpace;

    iter getChildIndices() : int {
      for idx in this.myLocaleSpace do
        yield idx;
    }

    proc getChild(idx:int) return this.myLocales[idx];

    iter getChldren() : locale  {
      for loc in this.myLocales do
        yield loc;
    }

    proc getDefaultLocaleSpace() return this.myLocaleSpace;
    proc getDefaultLocaleArray() return myLocales;

    proc localeIDtoLocale(id : chpl_localeID_t) {
      // In the default architecture, there are only nodes and no sublocales.
      // What is more, the nodeID portion of a wide pointer is the same as
      // the index into myLocales that yields the locale representing that node.
      return myLocales[id.node];
    }
  }

  if (rootLocale == nil) then
    rootLocale = new DefaultRootLocale();

  // Expose the underlying locales array (and its domain)
  // for user convenience and backward compatibility.
  // The downcase is because we cannot move the domain and array return types
  // into the base (abstract) locale class.
  // That would make locales depend on arrays which depend on locales....
  // If we had a way to express domain and array return types abstractly,
  // that problem would go away. <hilde>
  const Locales => (rootLocale:RootLocale).getDefaultLocaleArray();
  const LocaleSpace = Locales.domain;
}


Some cleanup

Some cleanup

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

// Dimension constanst
param NDIM=3;
param NEXT=2; // w, r2
param NTOT=NDIM+NEXT;
const W=NDIM;
const R2=NDIM+1;

// 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() {
  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, _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() {
  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);
}