From 659af4cd0f1d53dd6590acb1ad6babe57525afc7 Mon Sep 17 00:00:00 2001
From: ivan-aksamentov <ivan.aksamentov@gmail.com>
Date: Tue, 17 Jan 2023 08:25:42 +0100
Subject: [PATCH] feat: scale tree branches

---
 src/components/Tree/PhyloGraph/graph.ts | 10 ++++++----
 src/components/Tree/SpeciesTree.tsx     |  5 ++++-
 src/state/tree.state.ts                 |  9 +++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/components/Tree/PhyloGraph/graph.ts b/src/components/Tree/PhyloGraph/graph.ts
index d5cb613..badd0cc 100644
--- a/src/components/Tree/PhyloGraph/graph.ts
+++ b/src/components/Tree/PhyloGraph/graph.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-loops/no-loops */
-import { max, min, cloneDeep, sumBy, meanBy, isEmpty, last } from 'lodash'
+import { max, min, cloneDeep, sumBy, meanBy, isEmpty, last, isNil } from 'lodash'
 import { ErrorInternal } from 'src/helpers/ErrorInternal'
 
 import { PHYLO_GRAPH_NODE_RADIUS } from './constants'
@@ -77,6 +77,7 @@ export interface GraphEdge {
 
 export interface GraphLayoutOptions {
   mirrored?: boolean
+  scaleBranches?: boolean
 }
 
 export function calculateGraphLayout(
@@ -115,9 +116,10 @@ export function calculateGraphLayout(
       node.layout.maxDepth = 0
       node.layout.minDepth = 0
     } else {
-      node.layout.meanDepth = meanBy(parents, ([parent]) => parent.layout.meanDepth) + 1
-      node.layout.minDepth = (min(parents.map(([parent, _]) => parent.layout.meanDepth)) ?? 0) + 1
-      node.layout.maxDepth = (max(parents.map(([parent, _]) => parent.layout.meanDepth)) ?? 0) + 1
+      const depth = options?.scaleBranches && !isNil(node.branch_length) ? node.branch_length : 1
+      node.layout.meanDepth = meanBy(parents, ([parent]) => parent.layout.meanDepth) + depth
+      node.layout.minDepth = (min(parents.map(([parent, _]) => parent.layout.meanDepth)) ?? 0) + depth
+      node.layout.maxDepth = (max(parents.map(([parent, _]) => parent.layout.meanDepth)) ?? 0) + depth
     }
     depth = Math.max(depth, node.layout.maxDepth)
     return node
diff --git a/src/components/Tree/SpeciesTree.tsx b/src/components/Tree/SpeciesTree.tsx
index ec04ef2..52939e4 100644
--- a/src/components/Tree/SpeciesTree.tsx
+++ b/src/components/Tree/SpeciesTree.tsx
@@ -1,15 +1,18 @@
 import React, { memo, useMemo } from 'react'
+import { useRecoilValue } from 'recoil'
 import { useResizeDetector } from 'react-resize-detector'
 import type { SpeciesDesc } from 'src/hooks/useDataIndexQuery'
 import { useSpeciesTreeJson } from 'src/hooks/useDataIndexQuery'
 import { convertPhyloTreeToGraph } from 'src/components/Tree/PhyloGraph/convertPhyloTreeToGraph'
 import { PhyloGraph } from 'src/components/Tree/PhyloGraph/PhyloGraph'
+import { speciesTreeOptionsAtom } from 'src/state/tree.state'
 
 export interface SpeciesTreeProps {
   species: SpeciesDesc
 }
 
 function SpeciesTreeUnmemo({ species }: SpeciesTreeProps) {
+  const speciesTreeOptions = useRecoilValue(speciesTreeOptionsAtom)
   const { tree, meta } = useSpeciesTreeJson(species.id)
   const graph = useMemo(() => convertPhyloTreeToGraph(tree, meta), [meta, tree])
 
@@ -24,7 +27,7 @@ function SpeciesTreeUnmemo({ species }: SpeciesTreeProps) {
 
   return (
     <div className="w-100 h-100" ref={containerRef}>
-      <PhyloGraph width={width} height={height} graph={graph} />
+      <PhyloGraph width={width} height={height} graph={graph} options={speciesTreeOptions} />
     </div>
   )
 }
diff --git a/src/state/tree.state.ts b/src/state/tree.state.ts
index 3bf4a48..1f6dee6 100644
--- a/src/state/tree.state.ts
+++ b/src/state/tree.state.ts
@@ -6,9 +6,18 @@ export const highlightedNodeAtom = atomFamily<boolean, string>({
   default: false,
 })
 
+export const speciesTreeOptionsAtom = atom<GraphLayoutOptions>({
+  key: 'speciesTreeOptionsAtom',
+  default: {
+    mirrored: false,
+    scaleBranches: true,
+  },
+})
+
 export const geneTreeOptionsAtom = atom<GraphLayoutOptions>({
   key: 'geneTreeOptionsAtom',
   default: {
     mirrored: true,
+    scaleBranches: true,
   },
 })