Dev focus pca (#624)

* viewer camera change based on Pca

* minor code refactor

* update author

* Update src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* Update src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* Update src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* Update src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* Update src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* Update src/mol-plugin-ui/structure/components.tsx

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* Update src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

* revise

* deepclone

* chunked-array

---------

Co-authored-by: David Sehnal <dsehnal@users.noreply.github.com>

CHANGELOG.md

@@ -6,6 +6,9 @@ Note that since we don't clearly distinguish between a public and private interf
 
 ## [Unreleased]
 
+- Change the position of the camera based on the PCA of the structure and the following rules.
+    - The first residue should be in first quadrant if there is only one chain
+    - The average position of the residues of the first chain should be in the first residue if there are more than one chain.
 - Add `HeadlessPluginContext` and `HeadlessScreenshotHelper` to be used in Node.js
 - Add example `image-renderer`
 - Fix wrong offset when rendering text with orthographic projection

package.json

@@ -93,6 +93,7 @@
     "Adam Midlik <midlik@gmail.com>",
     "Koya Sakuma <koya.sakuma.work@gmail.com>",
     "Gianluca Tomasello <giagitom@gmail.com>",
+    "Ke Ma <mark.ma@rcsb.org>",
     "Jason Pattle <jpattle@exscientia.co.uk>",
     "David Williams <dwilliams@nobiastx.com>"
   ],

src/mol-plugin-state/manager/camera.ts

@@ -3,6 +3,7 @@
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Ke Ma <mark.ma@rcsb.org>
  */
 
 import { Sphere3D } from '../../mol-math/geometry';
@@ -13,6 +14,8 @@ import { Loci } from '../../mol-model/loci';
 import { BoundaryHelper } from '../../mol-math/geometry/boundary-helper';
 import { GraphicsRenderObject } from '../../mol-gl/render-object';
 import { StructureElement } from '../../mol-model/structure';
+import { Vec3 } from '../../mol-math/linear-algebra/3d/vec3';
+import { pcaFocus } from './focus-camera/focus-first-residue';
 
 // TODO: make this customizable somewhere?
 const DefaultCameraFocusOptions = {
@@ -84,7 +87,7 @@ export class CameraManager {
         }
     }
 
-    focusSpheres<T>(xs: ReadonlyArray<T>, sphere: (t: T) => Sphere3D | undefined, options?: Partial<CameraFocusOptions>) {
+    focusSpheres<T>(xs: ReadonlyArray<T>, sphere: (t: T) => Sphere3D | undefined, options?: Partial<CameraFocusOptions> & { principalAxes?: PrincipalAxes, positionToFlip?: Vec3 }) {
         const spheres = [];
 
         for (const x of xs) {
@@ -106,7 +109,7 @@ export class CameraManager {
         this.focusSphere(this.boundaryHelper.getSphere(), options);
     }
 
-    focusSphere(sphere: Sphere3D, options?: Partial<CameraFocusOptions> & { principalAxes?: PrincipalAxes }) {
+    focusSphere(sphere: Sphere3D, options?: Partial<CameraFocusOptions> & { principalAxes?: PrincipalAxes, positionToFlip?: Vec3 }) {
         const { canvas3d } = this.plugin;
         if (!canvas3d) return;
 
@@ -114,9 +117,10 @@ export class CameraManager {
         const radius = Math.max(sphere.radius + extraRadius, minRadius);
 
         if (options?.principalAxes) {
-            const { origin, dirA, dirC } = options?.principalAxes.boxAxes;
-            const snapshot = canvas3d.camera.getFocus(origin, radius, dirA, dirC);
-            canvas3d.requestCameraReset({ durationMs, snapshot });
+            this.plugin.canvas3d?.camera.setState(Camera.createDefaultSnapshot());
+            const { origin, dirA, dirC } = pcaFocus(this.plugin, options);
+            const snapshot = this.plugin.canvas3d?.camera.getFocus(origin, radius, dirA, dirC);
+            this.plugin.canvas3d?.requestCameraReset({ durationMs, snapshot });
         } else {
             const snapshot = canvas3d.camera.getFocus(sphere.center, radius);
             canvas3d.requestCameraReset({ durationMs, snapshot });

src/mol-plugin-state/manager/focus-camera/focus-first-residue.ts

@@ -0,0 +1,147 @@
+/**
+ * Copyright (c) 2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author Ke Ma <mark.ma@rcsb.org>
+ */
+import { Structure } from '../../../mol-model/structure';
+import { Vec3 } from '../../..//mol-math/linear-algebra/3d/vec3';
+import { PluginContext } from '../../../mol-plugin/context';
+import { CameraFocusOptions } from '../camera';
+import { PrincipalAxes } from '../../../mol-math/linear-algebra/matrix/principal-axes';
+import { StructureComponentRef } from '../structure/hierarchy-state';
+import { deepClone } from '../../../mol-util/object';
+import { ChunkedArray } from '../../../mol-data/util/chunked-array';
+
+
+export function getPolymerPositions(polymerStructure: Structure): Float32Array {
+    const tmpMatrix = Vec3.zero();
+    const cAdd3 = ChunkedArray.add3;
+    const positions = ChunkedArray.create(Float32Array, 3, 1024, polymerStructure.atomicResidueCount);
+    for (let i = 0; i < polymerStructure.units.length; i++) {
+        const unit = polymerStructure.units[i];
+        const { polymerElements } = unit.props;
+        const readPosition = unit.conformation.position;
+        if (polymerElements) {
+            for (let j = 0; j < polymerElements.length; j++) {
+                readPosition(polymerElements[j], tmpMatrix);
+                cAdd3(positions, tmpMatrix[0], tmpMatrix[1], tmpMatrix[2]);
+            }
+        }
+    }
+    return ChunkedArray.compact(positions) as Float32Array;
+}
+export function calculateDisplacement(position: Vec3, origin: Vec3, normalDir: Vec3) {
+    const A = normalDir[0];
+    const B = normalDir[1];
+    const C = normalDir[2];
+    const D = -A * origin[0] - B * origin[1] - C * origin[2];
+
+    const x = position[0];
+    const y = position[1];
+    const z = position[2];
+
+    const displacement = (A * x + B * y + C * z + D) / Math.sqrt(A * A + B * B + C * C);
+    return displacement;
+}
+
+export function getAxesToFlip(position: Vec3, origin: Vec3, up: Vec3, normalDir: Vec3) {
+    const toYAxis = calculateDisplacement(position, origin, normalDir);
+    const toXAxis = calculateDisplacement(position, origin, up);
+    const axes: ('aroundX' | 'aroundY')[] = [];
+    if (toYAxis < 0) axes.push('aroundY');
+    if (toXAxis < 0) axes.push('aroundX');
+    return axes;
+}
+
+export function getFirstResidueOrAveragePosition(structure: Structure, caPositions: Float32Array): Vec3 {
+    if (structure.units.length === 1) {
+        // if only one chain => first residue coordinates
+        return Vec3.create(caPositions[0], caPositions[1], caPositions[2]);
+    } else {
+        // if more than one chain => average of coordinates of the first chain
+        const tmpMatrixPos = Vec3.zero();
+        const atomIndices = structure.units[0].props.polymerElements;
+        const firstChainPositions = [];
+        if (atomIndices) {
+            for (let i = 0; i < atomIndices.length; i++) {
+                const coordinates = structure.units[0].conformation.position(atomIndices[i], tmpMatrixPos);
+                for (let j = 0; j < coordinates.length; j++) {
+                    firstChainPositions.push(coordinates[j]);
+                }
+            }
+            let sumX = 0;
+            let sumY = 0;
+            let sumZ = 0;
+            for (let i = 0; i < firstChainPositions.length; i += 3) {
+                sumX += firstChainPositions[i];
+                sumY += firstChainPositions[i + 1];
+                sumZ += firstChainPositions[i + 2];
+            }
+            const averagePosition = Vec3.zero();
+            averagePosition[0] = sumX / atomIndices.length;
+            averagePosition[1] = sumY / atomIndices.length;
+            averagePosition[2] = sumZ / atomIndices.length;
+            return averagePosition;
+        } else {
+            return Vec3.create(caPositions[0], caPositions[1], caPositions[2]);
+        }
+    }
+
+}
+
+export function pcaFocus(plugin: PluginContext, options: Partial<CameraFocusOptions> & { principalAxes?: PrincipalAxes, positionToFlip?: Vec3 }) {
+    if (options?.principalAxes) {
+        const { origin, dirA, dirB, dirC } = options.principalAxes.boxAxes;
+        let toFlip: ('aroundX' | 'aroundY')[] = [];
+        if (options.positionToFlip) {
+            toFlip = getAxesToFlip(options.positionToFlip, origin, dirA, dirB);
+        }
+        toFlip.forEach((axis)=>{
+            if (axis === 'aroundY') {
+                Vec3.negate(dirC, dirC);
+            } else if (axis === 'aroundX') {
+                Vec3.negate(dirA, dirA);
+                Vec3.negate(dirC, dirC);
+            }
+        });
+        if (plugin.canvas3d) {
+            const position = Vec3();
+            Vec3.scaleAndAdd(position, position, origin, 100);
+            plugin.canvas3d.camera.setState({ position }, 0);
+            const deltaDistance = Vec3();
+            Vec3.negate(deltaDistance, position);
+            if (Vec3.dot(deltaDistance, dirC) <= 0) {
+                Vec3.negate(plugin.canvas3d.camera.position, position);
+            }
+            const up = Vec3.create(0, 1, 0);
+            if (Vec3.dot(up, dirA) <= 0) {
+                Vec3.negate(plugin.canvas3d?.camera.up, plugin.canvas3d.camera.up);
+            }
+        }
+        return { origin, dirA, dirB, dirC };
+    }
+    return {
+        origin: Vec3.zero(),
+        dirA: Vec3.zero(),
+        dirB: Vec3.zero(),
+        dirC: Vec3.zero()
+    };
+}
+
+export function getPcaTransform(group: StructureComponentRef[]): { principalAxes?: PrincipalAxes, positionToFlip?: Vec3 } | undefined {
+    const polymerStructure = group[0].cell.obj?.data;
+    if (!polymerStructure) {
+        return undefined;
+    }
+    if ('_pcaTransformData' in polymerStructure.currentPropertyData) {
+        return deepClone(polymerStructure.currentPropertyData._pcaTransformData);
+    }
+    if (!polymerStructure.units[0]?.props.polymerElements?.length) {
+        polymerStructure.currentPropertyData._pcaTransformData = undefined;
+        return undefined;
+    }
+    const positions = getPolymerPositions(polymerStructure);
+    const positionToFlip = getFirstResidueOrAveragePosition(polymerStructure, positions);
+    polymerStructure.currentPropertyData._pcaTransformData = { principalAxes: PrincipalAxes.ofPositions(positions), positionToFlip };
+    return deepClone(polymerStructure.currentPropertyData._pcaTransformData);
+}

src/mol-plugin-ui/structure/components.tsx

@@ -1,12 +1,14 @@
 /**
- * Copyright (c) 2020 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ * Copyright (c) 2020-2022 mol* contributors, licensed under MIT, See LICENSE file for more info.
  *
  * @author David Sehnal <david.sehnal@gmail.com>
  * @author Alexander Rose <alexander.rose@weirdbyte.de>
+ * @author Ke Ma <mark.ma@rcsb.org>
  */
 
 import * as React from 'react';
 import { getStructureThemeTypes } from '../../mol-plugin-state/helpers/structure-representation-params';
+import { getPcaTransform } from '../../mol-plugin-state/manager/focus-camera/focus-first-residue';
 import { StructureComponentManager } from '../../mol-plugin-state/manager/structure/component';
 import { StructureHierarchyManager } from '../../mol-plugin-state/manager/structure/hierarchy';
 import { StructureComponentRef, StructureRepresentationRef } from '../../mol-plugin-state/manager/structure/hierarchy-state';
@@ -315,7 +317,7 @@ class StructureComponentGroup extends PurePluginUIComponent<{ group: StructureCo
         this.plugin.managers.camera.focusSpheres(this.props.group, e => {
             if (e.cell.state.isHidden) return;
             return e.cell.obj?.data.boundary.sphere;
-        });
+        }, getPcaTransform(this.props.group));
     };
 
     get reprLabel() {