Merge pull request #1484 from nextstrain/fix/regression

Regression line only considers visible nodes

src/components/tree/phyloTree/change.js

@@ -333,6 +333,11 @@ export const change = function change({
   if (changeColorBy) {
     this.updateColorBy();
   }
+  // recalculate existing regression if needed
+  if (changeVisibility && this.regression) {
+    elemsToUpdate.add(".regression");
+    this.calculateRegression(); // Note: must come after `updateNodesWithNewData()`
+  }
   /* some things need to update d.inView and/or d.update. This should be centralised */
   /* TODO: list all functions which modify these */
   if (zoomIntoClade) { /* must happen below updateNodesWithNewData */

src/components/tree/phyloTree/layouts.js

@@ -3,7 +3,6 @@
 import { min, max } from "d3-array";
 import scaleLinear from "d3-scale/src/linear";
 import {point as scalePoint} from "d3-scale/src/band";
-import { calculateRegressionThroughRoot, calculateRegressionWithFreeIntercept } from "./regression";
 import { timerStart, timerEnd } from "../../../util/perf";
 import { getTraitFromNode, getDivFromNode } from "../../../util/treeMiscHelpers";
 
@@ -126,11 +125,7 @@ export const scatterplotLayout = function scatterplotLayout() {
   }
 
   if (this.scatterVariables.showRegression) {
-    if (this.layout==="clock") {
-      this.regression = calculateRegressionThroughRoot(this.nodes);
-    } else {
-      this.regression = calculateRegressionWithFreeIntercept(this.nodes);
-    }
+    this.calculateRegression(); // sets this.regression
   }
 
 };

src/components/tree/phyloTree/phyloTree.js

@@ -7,6 +7,7 @@ import * as layouts from "./layouts";
 import * as grid from "./grid";
 import * as confidence from "./confidence";
 import * as labels from "./labels";
+import * as regression from "./regression";
 
 /* phylogenetic tree drawing function - the actual tree is rendered by the render prototype */
 const PhyloTree = function PhyloTree(reduxNodes, id, idxOfInViewRootNode) {
@@ -68,6 +69,7 @@ PhyloTree.prototype.unrootedLayout = layouts.unrootedLayout;
 PhyloTree.prototype.radialLayout = layouts.radialLayout;
 PhyloTree.prototype.setScales = layouts.setScales;
 PhyloTree.prototype.mapToScreen = layouts.mapToScreen;
+PhyloTree.prototype.calculateRegression = regression.calculateRegression;
 
 /* C O N F I D E N C E    I N T E R V A L S */
 PhyloTree.prototype.removeConfidence = confidence.removeConfidence;

src/components/tree/phyloTree/regression.js

@@ -1,16 +1,19 @@
 import { sum } from "d3-array";
 import { formatDivergence, guessAreMutationsPerSite} from "./helpers";
+import { NODE_VISIBLE } from "../../../util/globals";
 
 
 /**
  * this function calculates a regression between
- * the x and y values of terminal nodes, passing through
- * nodes[0].
- * It does not consider which tips are inView / visible.
+ * the x and y values of terminal nodes which are also visible.
+ * The regression is forced to pass through nodes[0].
  */
-export function calculateRegressionThroughRoot(nodes) {
-  const terminalNodes = nodes.filter((d) => !d.n.hasChildren);
+function calculateRegressionThroughRoot(nodes) {
+  const terminalNodes = nodes.filter((d) => !d.n.hasChildren && d.visibility === NODE_VISIBLE);
   const nTips = terminalNodes.length;
+  if (nTips===0) {
+    return {slope: undefined, intercept: undefined, r2: undefined};
+  }
   const offset = nodes[0].x;
   const XY = sum(
     terminalNodes.map((d) => (d.y) * (d.x - offset))
@@ -25,13 +28,17 @@ export function calculateRegressionThroughRoot(nodes) {
 }
 
 /**
- * Calculate regression through terminal nodes which have both x & y values
+ * Calculate regression through visible terminal nodes which have both x & y values
  * set. These values must be numeric.
- * This function does not consider which tips are inView / visible.
  */
-export function calculateRegressionWithFreeIntercept(nodes) {
-  const terminalNodesWithXY = nodes.filter((d) => (!d.n.hasChildren) && d.x!==undefined && d.y!==undefined);
+function calculateRegressionWithFreeIntercept(nodes) {
+  const terminalNodesWithXY = nodes.filter(
+    (d) => (!d.n.hasChildren) && d.x!==undefined && d.y!==undefined && d.visibility === NODE_VISIBLE
+  );
   const nTips = terminalNodesWithXY.length;
+  if (nTips===0) {
+    return {slope: undefined, intercept: undefined, r2: undefined};
+  }
   const meanX = sum(terminalNodesWithXY.map((d) => d.x))/nTips;
   const meanY = sum(terminalNodesWithXY.map((d) => d.y))/nTips;
   const slope = sum(terminalNodesWithXY.map((d) => (d.x-meanX)*(d.y-meanY))) /
@@ -42,6 +49,14 @@ export function calculateRegressionWithFreeIntercept(nodes) {
   return {slope, intercept, r2};
 }
 
+/** sets this.regression  */
+export function calculateRegression() {
+  if (this.layout==="clock") {
+    this.regression = calculateRegressionThroughRoot(this.nodes);
+  } else {
+    this.regression = calculateRegressionWithFreeIntercept(this.nodes);
+  }
+}
 
 export function makeRegressionText(regression, layout, yScale) {
   if (layout==="clock") {

src/components/tree/phyloTree/renderers.js

@@ -241,6 +241,11 @@ export const drawBranches = function drawBranches() {
  * @return {null}
  */
 export const drawRegression = function drawRegression() {
+  /* check we have computed a sensible regression before attempting to draw */
+  if (this.regression.slope===undefined) {
+    return;
+  }
+
   const leftY = this.yScale(this.regression.intercept + this.xScale.domain()[0] * this.regression.slope);
   const rightY = this.yScale(this.regression.intercept + this.xScale.domain()[1] * this.regression.slope);