Color the 3D structure by a property

python/examples/colors.py

@@ -0,0 +1,70 @@
+"""
+Atom property coloring
+======================
+
+This example demonstrates how to color atoms based on scalar properties.
+
+Note that the same parameters can be used with `chemiscope.show` to visualize an
+interactive widget in a Jupyter notebook.
+"""
+import ase.io
+import numpy as np
+
+import chemiscope
+
+# loads a dataset of structures
+frames = ase.io.read("data/alpha-mu.xyz", ":")
+
+# compute some scalar quantities to display as atom coloring
+polarizability = []
+alpha_eigenvalues = []
+anisotropy = []
+
+for frame in frames:
+    for axx, ayy, azz, axy, axz, ayz in zip(
+        frame.arrays["axx"],
+        frame.arrays["ayy"],
+        frame.arrays["azz"],
+        frame.arrays["axy"],
+        frame.arrays["axz"],
+        frame.arrays["ayz"],
+    ):
+        polarizability.append((axx + ayy + azz) / 3)
+
+        # one possible measure of anisotropy...
+        eigenvalues = np.linalg.eigvalsh(
+            [[axx, axy, axz], [axy, ayy, ayz], [axz, ayz, azz]]
+        )
+        alpha_eigenvalues.append(eigenvalues)
+
+        anisotropy.append(eigenvalues[2] - eigenvalues[0])
+
+
+# now we just write the chemiscope input file
+chemiscope.write_input(
+    "colors-example.json.gz",
+    frames=frames,
+    # properties can be extracted from the ASE.Atoms frames
+    properties={
+        "polarizability": np.vstack(polarizability),
+        "anisotropy": np.vstack(anisotropy),
+        "alpha_eigenvalues": np.vstack(alpha_eigenvalues),
+    },
+    # the write_input function also allows defining the default visualization settings
+    settings={
+        "map": {
+            "x": {"property": "alpha_eigenvalues[1]"},
+            "y": {"property": "alpha_eigenvalues[2]"},
+            "z": {"property": "alpha_eigenvalues[3]"},
+            "color": {"property": "anisotropy"},
+        },
+        "structure": [
+            {
+                "color": {"property": "anisotropy"},
+            }
+        ],
+    },
+    # the properties we want to visualise are atomic properties - in order to view them
+    # in the map panel, we must indicate that all atoms are environment centers
+    environments=chemiscope.all_atomic_environments(frames),
+)

src/index.ts

@@ -190,6 +190,7 @@ class DefaultVisualizer {
             config.structure,
             this._indexer,
             dataset.structures,
+            dataset.properties,
             dataset.environments,
             config.maxStructureViewers
         );
@@ -484,6 +485,7 @@ class StructureVisualizer {
             config.structure,
             this._indexer,
             dataset.structures,
+            dataset.properties,
             dataset.environments,
             1 // only allow one structure
         );

src/info/info.ts

@@ -18,7 +18,7 @@ import { Table } from './table';
 import INFO_SVG from '../static/info.svg';
 import * as styles from '../styles';
 
-function filter<T extends Record<string, Property>>(
+export function filter<T extends Record<string, Property>>(
     obj: T,
     predicate: (o: Property) => boolean
 ): Record<string, Property> {

src/static/chemiscope.css

@@ -500,3 +500,32 @@
     margin-top: 3px;
     margin-bottom: 1px;
 }
+
+.chsp-atom-color-property {
+    display: grid;
+    align-items: center;
+    justify-items: center;
+    grid-template-columns: 1fr 1fr;
+    column-gap: 1em;
+    margin: 0.5em 0;
+}
+
+.chsp-atom-extra-options {
+    display: grid;
+    align-items: center;
+    justify-items: center;
+    grid-template-columns: auto 1fr 1fr;
+    column-gap: 1em;
+}
+
+.chsp-atom-extra-options > * {
+    margin: 0.5em 0;
+}
+
+.chsp-atom-extra-options label {
+    min-width: auto;
+}
+
+.chsp-atom-extra-options {
+    grid-column: 1 / span 2;
+}

src/structure/grid.ts

@@ -8,6 +8,7 @@ import assert from 'assert';
 import {
     Environment,
     JsObject,
+    Property,
     Settings,
     Structure,
     UserStructure,
@@ -21,10 +22,17 @@ import { enumerate, generateGUID, getByID, getFirstKey, getNextColor, sendWarnin
 
 import { LoadOptions, MoleculeViewer } from './viewer';
 
+import { filter } from '../info/info';
+
 import CLOSE_SVG from '../static/close.svg';
 import DUPLICATE_SVG from '../static/duplicate.svg';
 import PNG_SVG from '../static/download-png.svg';
 
+/// Extension of `Environment` adding the global index of the environment
+interface NumberedEnvironment extends Environment {
+    index: number;
+}
+
 /**
  * Create a list of environments grouped together by structure.
  *
@@ -41,7 +49,7 @@ import PNG_SVG from '../static/download-png.svg';
 function groupByStructure(
     structures: (Structure | UserStructure)[],
     environments?: Environment[]
-): Environment[][] | undefined {
+): NumberedEnvironment[][] | undefined {
     if (environments === undefined) {
         return undefined;
     }
@@ -50,11 +58,15 @@ function groupByStructure(
         Array.from({ length: structures[i].size })
     );
 
-    for (const env of environments) {
-        result[env.structure][env.center] = env;
+    for (let i = 0; i < environments.length; i++) {
+        const env = environments[i];
+        result[env.structure][env.center] = {
+            index: i,
+            ...env,
+        };
     }
 
-    return result as Environment[][];
+    return result as NumberedEnvironment[][];
 }
 
 interface ViewerGridData {
@@ -116,10 +128,12 @@ export class ViewersGrid {
     private _root: HTMLElement;
     /// List of structures in the dataset
     private _structures: Structure[] | UserStructure[];
+    /// List of per-atom properties in the dataset
+    private _properties: Record<string, number[]>;
     /// Cached string representation of structures
     private _resolvedStructures: Structure[];
     /// Optional list of environments for each structure
-    private _environments?: Environment[][];
+    private _environments?: NumberedEnvironment[][];
     /// Maximum number of allowed structure viewers
     private _maxViewers: number;
     /// The indexer translating between environments indexes and structure/atom
@@ -153,10 +167,27 @@ export class ViewersGrid {
         element: string | HTMLElement,
         indexer: EnvironmentIndexer,
         structures: Structure[] | UserStructure[],
+        properties?: { [name: string]: Property },
         environments?: Environment[],
         maxViewers: number = 9
     ) {
         this._structures = structures;
+        if (properties === undefined) {
+            this._properties = {};
+        } else {
+            const numberProperties = filter(properties, (p) =>
+                Object.values(p.values).every((v) => typeof v === 'number')
+            );
+            const atomProperties = filter(numberProperties, (p) => p.target === 'atom');
+
+            const props: Record<string, number[]> = Object.fromEntries(
+                Object.entries(atomProperties).map(([key, value]) => [
+                    key,
+                    value.values as number[],
+                ])
+            );
+            this._properties = props;
+        }
         this._resolvedStructures = new Array<Structure>(structures.length);
         this._environments = groupByStructure(this._structures, environments);
         this._indexer = indexer;
@@ -478,6 +509,38 @@ export class ViewersGrid {
         return this._resolvedStructures[index];
     }
 
+    /**
+     * Get the values of the properties for all the atoms in the current structure
+     *
+     * @param structure index of the current structure
+     * @returns
+     */
+    private _propertiesForStructure(
+        structure: number
+    ): Record<string, (number | undefined)[]> | undefined {
+        const properties: Record<string, (number | undefined)[]> = {};
+
+        if (this._environments !== undefined) {
+            const environments = this._environments[structure];
+            for (const name in this._properties) {
+                const values = this._properties[name];
+
+                properties[name] = [];
+                for (const environment of environments) {
+                    if (environment !== undefined) {
+                        properties[name].push(values[environment.index]);
+                    } else {
+                        properties[name].push(undefined);
+                    }
+                }
+            }
+
+            return properties;
+        } else {
+            return undefined;
+        }
+    }
+
     private _showInViewer(guid: GUID, indexes: Indexes): void {
         const data = this._cellsData.get(guid);
         assert(data !== undefined);
@@ -496,7 +559,11 @@ export class ViewersGrid {
                 }
             }
 
-            viewer.load(this._structure(indexes.structure), options);
+            viewer.load(
+                this._structure(indexes.structure),
+                this._propertiesForStructure(indexes.structure),
+                options
+            );
             data.current = indexes;
         }
 
@@ -695,7 +762,11 @@ export class ViewersGrid {
 
             // add a new cells if necessary
             if (!this._cellsData.has(cellGUID)) {
-                const viewer = new MoleculeViewer(this._getById<HTMLElement>(`gi-${cellGUID}`));
+                const propertiesName = this._properties ? Object.keys(this._properties) : [];
+                const viewer = new MoleculeViewer(
+                    this._getById<HTMLElement>(`gi-${cellGUID}`),
+                    propertiesName
+                );
 
                 viewer.onselect = (atom: number) => {
                     if (this._indexer.mode !== 'atom' || this._active !== cellGUID) {

src/structure/options.html.in

@@ -33,6 +33,59 @@
                   <label class="form-check-label" for="rotation" title="rotate the camera continuously">rotation</label>
                </div>
             </div>
+            <div class="chsp-atom-color-property">
+
+               <div class="input-group input-group-sm">
+                  <label class="input-group-text" for="atom-color-property">color:</label>
+                  <select id="atom-color-property" class="form-select">
+                     <option value="element">element</option>
+                  </select>
+               </div>
+
+               <button
+                  id="atom-color-more-options"
+                  class="btn btn-sm btn-secondary chsp-extra-options-btn"
+                  type="button"
+                  data-bs-toggle="collapse"
+                  data-bs-target="#atom-color-extra"
+               >
+                  more options
+               </button>
+
+               <div class="collapse" id="atom-color-extra" style="grid-column: auto / span 2; margin-top: 0.3em">
+                  <div class="chsp-atom-extra-options">
+                     <div class="input-group input-group-sm">
+                        <label class="input-group-text" for="atom-color-transform">transform:</label>
+                        <select id="atom-color-transform" class="form-select" disabled>
+                           <option value="linear">linear</option>
+                           <option value="log">log10</option>
+                           <option value="sqrt">sqrt</option>
+                           <option value="inverse">inverse</option>
+                        </select>
+                     </div>
+
+                     <div class="input-group input-group-sm">
+                        <label id="atom-color-min-label" class="input-group-text" for="atom-color-min"> min: </label>
+                        <input id="atom-color-min" class="form-control" type="number" step="any" />
+                     </div>
+
+                     <div class="input-group input-group-sm">
+                        <label id="atom-color-max-label" class="input-group-text" for="atom-color-max"> max: </label>
+                        <input id="atom-color-max" class="form-control" type="number" step="any" />
+                     </div>
+                  </div>
+
+                  <div class="chsp-atom-extra-options" style="grid-template-columns: 1fr 1fr">
+                     <div class="input-group input-group-sm">
+                        <label class="input-group-text" for="atom-color-palette"> palette: </label>
+                        <select id="atom-color-palette" class="form-select"></select>
+                     </div>
+
+                     <button class="btn btn-sm btn-secondary" type="button" id="atom-color-reset" style="width: 100%">reset min/max</button>
+                  </div>
+               </div>
+
+            </div>
             <div class="chsp-hide-if-no-shapes" style="margin-top: 1em">
                <div class="input-group input-group-sm">
                   <label
@@ -74,6 +127,7 @@
                   <option class="chsp-hide-if-no-cell" value="abc">abc</option>
                </select>
             </div>
+                <div style="margin-bottom: 0.8em"></div>
             <div class="chsp-hide-if-no-environments">
                <div style="margin-bottom: 1.5em"></div>
                <h5 class="chsp-settings-section-title">Environments</h5>
@@ -115,6 +169,7 @@
                      <div class="input-group input-group-sm">
                         <label class="input-group-text" for="env-bg-color">color</label>
                         <select id="env-bg-color" class="form-select">
+                           <option value="prop">prop</option>
                            <option value="CPK">CPK</option>
                            <option value="grey">grey</option>
                         </select>