diff --git a/python/examples/colors.py b/python/examples/colors.py new file mode 100644 index 000000000..1e37f2f2a --- /dev/null +++ b/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), +) diff --git a/src/index.ts b/src/index.ts index 87d0a8149..1aa81fd8a 100644 --- a/src/index.ts +++ b/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 ); diff --git a/src/info/info.ts b/src/info/info.ts index 16bac9e67..95b5b3899 100644 --- a/src/info/info.ts +++ b/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>( +export function filter>( obj: T, predicate: (o: Property) => boolean ): Record { diff --git a/src/static/chemiscope.css b/src/static/chemiscope.css index b6c0dea41..550e70c4c 100644 --- a/src/static/chemiscope.css +++ b/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; +} diff --git a/src/structure/grid.ts b/src/structure/grid.ts index 1a1dce88c..bd3befd0b 100644 --- a/src/structure/grid.ts +++ b/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; /// 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 = Object.fromEntries( + Object.entries(atomProperties).map(([key, value]) => [ + key, + value.values as number[], + ]) + ); + this._properties = props; + } this._resolvedStructures = new Array(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 | undefined { + const properties: Record = {}; + + 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(`gi-${cellGUID}`)); + const propertiesName = this._properties ? Object.keys(this._properties) : []; + const viewer = new MoleculeViewer( + this._getById(`gi-${cellGUID}`), + propertiesName + ); viewer.onselect = (atom: number) => { if (this._indexer.mode !== 'atom' || this._active !== cellGUID) { diff --git a/src/structure/options.html.in b/src/structure/options.html.in index 9f931ed6e..45927aa86 100644 --- a/src/structure/options.html.in +++ b/src/structure/options.html.in @@ -33,6 +33,59 @@ +
+ +
+ + +
+ + + +
+
+
+ + +
+ +
+ + +
+ +
+ + +
+
+ +
+
+ + +
+ + +
+
+ +
+
Environments
@@ -115,6 +169,7 @@
diff --git a/src/structure/options.ts b/src/structure/options.ts index 8e34da97f..f88402b64 100644 --- a/src/structure/options.ts +++ b/src/structure/options.ts @@ -51,6 +51,13 @@ export class StructureOptions extends OptionsGroup { // which colors for atoms not in the environment bgColor: HTMLOption<'string'>; }; + public color: { + property: HTMLOption<'string'>; + transform: HTMLOption<'string'>; + min: HTMLOption<'number'>; + max: HTMLOption<'number'>; + palette: HTMLOption<'string'>; + }; /// The Modal instance private _modal: Modal; @@ -59,7 +66,11 @@ export class StructureOptions extends OptionsGroup { // Callback to get the initial positioning of the settings modal. private _positionSettingsModal: PositioningCallback; - constructor(root: HTMLElement, positionSettings: PositioningCallback) { + constructor( + root: HTMLElement, + positionSettings: PositioningCallback, + propertiesName: string[] = [] + ) { super(); this.bonds = new HTMLOption('boolean', true); @@ -98,8 +109,31 @@ export class StructureOptions extends OptionsGroup { cutoff: new HTMLOption('number', 4.0), }; + this.color = { + property: new HTMLOption('string', 'element'), + min: new HTMLOption('number', 0), + max: new HTMLOption('number', 0), + transform: new HTMLOption('string', 'linear'), + palette: new HTMLOption('string', 'BWR'), + }; + + // validate atom properties for coloring + if (propertiesName.includes('element')) { + this.color.property.validate = optionValidator(propertiesName, 'color'); + } else { + this.color.property.validate = optionValidator( + propertiesName.concat(['element']), + 'color' + ); + } + this.color.transform.validate = optionValidator( + ['linear', 'log', 'sqrt', 'inverse'], + 'transform' + ); + this.color.palette.validate = optionValidator(['BWR', 'ROYGB', 'Sinebow'], 'palette'); + this.environments.bgColor.validate = optionValidator( - ['grey', 'CPK'], + ['grey', 'CPK', 'prop'], 'background atoms coloring' ); this.environments.bgStyle.validate = optionValidator( @@ -121,8 +155,7 @@ export class StructureOptions extends OptionsGroup { ).adoptedStyleSheets; this._openModal = openModal; root.appendChild(this._openModal); - - this._bind(); + this._bind(propertiesName); } /** Get in a element in the modal from its id */ @@ -223,7 +256,7 @@ export class StructureOptions extends OptionsGroup { } /** Bind all options to the corresponding HTML elements */ - private _bind(): void { + private _bind(propertiesName: string[]): void { this.atomLabels.bind(this.getModalElement('atom-labels'), 'checked'); const selectShape = this.getModalElement('shapes'); @@ -240,6 +273,29 @@ export class StructureOptions extends OptionsGroup { this.supercell[1].bind(this.getModalElement('supercell-b'), 'value'); this.supercell[2].bind(this.getModalElement('supercell-c'), 'value'); + // ======= data used as color values + const selectColorProperty = this.getModalElement('atom-color-property'); + // first option is 'element' + selectColorProperty.options.length = 0; + if (!propertiesName.includes('element')) { + selectColorProperty.options.add(new Option('element', 'element')); + } + for (const property of propertiesName) { + selectColorProperty.options.add(new Option(property, property)); + } + this.color.property.bind(selectColorProperty, 'value'); + this.color.transform.bind(this.getModalElement('atom-color-transform'), 'value'); + this.color.min.bind(this.getModalElement('atom-color-min'), 'value'); + this.color.max.bind(this.getModalElement('atom-color-max'), 'value'); + + // ======= color palette + const selectPalette = this.getModalElement('atom-color-palette'); + selectPalette.options.length = 0; + for (const key of ['BWR', 'ROYGB', 'Sinebow']) { + selectPalette.options.add(new Option(key, key)); + } + this.color.palette.bind(selectPalette, 'value'); + this.axes.bind(this.getModalElement('axes'), 'value'); this.keepOrientation.bind(this.getModalElement('keep-orientation'), 'checked'); this.playbackDelay.bind(this.getModalElement('playback-delay'), 'value'); diff --git a/src/structure/viewer.ts b/src/structure/viewer.ts index 0ee430987..faa2ea2cd 100644 --- a/src/structure/viewer.ts +++ b/src/structure/viewer.ts @@ -8,7 +8,7 @@ import assert from 'assert'; import * as $3Dmol from '3dmol'; import { assignBonds } from './assignBonds'; -import { getElement, unreachable } from '../utils'; +import { arrayMaxMin, getElement, sendWarning, unreachable } from '../utils'; import { PositioningCallback } from '../utils'; import { Environment, Settings, Structure } from '../dataset'; @@ -152,6 +152,12 @@ export class MoleculeViewer { }; /// List of atom-centered environments for the current structure private _environments?: (Environment | undefined)[]; + // List of properties for the current structure + private _properties?: Record; + // Button used to reset the range of color axis + private _colorReset: HTMLButtonElement; + // Button used to see more color options + private _colorMoreOptions: HTMLButtonElement; /** * Create a new `MoleculeViewer` inside the HTML DOM element with the given `id`. @@ -159,7 +165,7 @@ export class MoleculeViewer { * @param element HTML element or HTML id of the DOM element * where the viewer will be created */ - constructor(element: string | HTMLElement) { + constructor(element: string | HTMLElement, propertiesName: string[]) { const containerElement = getElement(element); const hostElement = document.createElement('div'); containerElement.appendChild(hostElement); @@ -216,10 +222,10 @@ export class MoleculeViewer { noShapeStyle, ]; - // Options reuse the same style sheets so they must be created after these. - - this._options = new StructureOptions(this._root, (rect) => - this.positionSettingsModal(rect) + this._options = new StructureOptions( + this._root, + (rect) => this.positionSettingsModal(rect), + propertiesName ); this._options.modal.shadow.adoptedStyleSheets = [ @@ -228,6 +234,9 @@ export class MoleculeViewer { noEnvsStyle, noShapeStyle, ]; + this._colorReset = this._options.getModalElement('atom-color-reset'); + this._colorMoreOptions = + this._options.getModalElement('atom-color-more-options'); this._connectOptions(); this._trajectoryOptions = this._options.getModalElement('trajectory-settings-group'); @@ -312,11 +321,17 @@ export class MoleculeViewer { * @param structure structure to load * @param options options for the new structure */ - public load(structure: Structure, options: Partial = {}): void { + public load( + structure: Structure, + properties?: Record | undefined, + options: Partial = {} + ): void { // if the canvas size changed since last structure, make sure we update // everything this.resize(); + this._properties = properties; + let previousDefaultCutoff = undefined; if (this._highlighted !== undefined) { previousDefaultCutoff = this._defaultCutoff(this._highlighted.center); @@ -779,6 +794,128 @@ export class MoleculeViewer { restyleAndRender(); }); + // ======= color settings + // setup state when the property changes + const colorPropertyChanged = () => { + const property = this._options.color.property.value; + + if (property !== 'element') { + this._options.color.transform.enable(); + this._options.color.transform.value = 'linear'; + this._options.color.min.enable(); + this._options.color.max.enable(); + + this._colorReset.disabled = false; + this._colorMoreOptions.disabled = false; + this._options.color.palette.enable(); + + const values = this._colorValues(property, 'linear'); + + if (values.some((v) => v === null)) { + sendWarning( + 'The selected structure has undefined properties for some atoms, these atoms will be colored in light gray.' + ); + } + + // To change min and max values when the transform has been changed + const { max, min } = arrayMaxMin(values); + + // We have to set max first and min second here to avoid sending + // a spurious warning in `colorRangeChange` below in case the + // new min is bigger than the old max. + this._options.color.min.value = Number.NEGATIVE_INFINITY; + this._options.color.max.value = max; + this._options.color.min.value = min; + this._setScaleStep([min, max]); + } else { + this._options.color.transform.disable(); + this._options.color.min.disable(); + this._options.color.max.disable(); + + this._colorReset.disabled = true; + this._colorMoreOptions.disabled = true; + this._options.color.palette.disable(); + + this._viewer.setColorByElement({}, $3Dmol.elementColors.Jmol); + } + restyleAndRender(); + }; + this._options.color.property.onchange.push(colorPropertyChanged); + + const colorRangeChange = (minOrMax: 'min' | 'max') => { + const min = this._options.color.min.value; + const max = this._options.color.max.value; + if (min > max) { + sendWarning( + `The inserted min and max values in color are such that min > max! The last inserted value was reset.` + ); + if (minOrMax === 'min') { + this._options.color.min.reset(); + } else { + this._options.color.max.reset(); + } + return; + } + this._setScaleStep([min, max]); + }; + + // ======= color transform + this._options.color.transform.onchange.push(() => { + const property = this._options.color.property.value; + assert(property !== 'element'); + const transform = this._options.color.transform.value; + + const values = this._colorValues(property, transform); + // To change min and max values when the transform has been changed + const { min, max } = arrayMaxMin(values); + + // to avoid sending a spurious warning in `colorRangeChange` below + // in case the new min is bigger than the old max. + this._options.color.min.value = Number.NEGATIVE_INFINITY; + this._options.color.max.value = max; + this._options.color.min.value = min; + this._setScaleStep([min, max]); + + restyleAndRender(); + }); + + // ======= color min + this._options.color.min.onchange.push(() => { + colorRangeChange('min'); + restyleAndRender(); + }); + + // ======= color max + this._options.color.max.onchange.push(() => { + colorRangeChange('max'); + restyleAndRender(); + }); + + // ======= color reset + this._colorReset.addEventListener('click', () => { + const properties = JSON.parse(JSON.stringify(this._properties)) as Record< + string, + (number | undefined)[] + >; + const property: string = this._options.color.property.value; + // Use map to extract the specified property values into an array + const values: number[] = properties[property].filter( + (value) => !isNaN(Number(value)) + ) as number[]; + // To change min and max values when the transform has been changed + const [min, max]: [number, number] = [Math.min(...values), Math.max(...values)]; + this._options.color.min.value = min; + this._options.color.max.value = max; + this._setScaleStep([min, max]); + + restyleAndRender(); + }); + + // ======= color palette + this._options.color.palette.onchange.push(() => { + restyleAndRender(); + }); + // Setup various buttons this._resetEnvCutoff = this._options.getModalElement('env-reset'); this._resetEnvCutoff.onclick = () => { @@ -1037,17 +1174,104 @@ export class MoleculeViewer { if (this._options.atoms.value) { style.sphere = { scale: this._options.spaceFilling.value ? 1.0 : 0.22, - }; + colorfunc: this._colorFunction(), + } as unknown as $3Dmol.SphereStyleSpec; } if (this._options.bonds.value) { style.stick = { radius: 0.15, - }; + colorfunc: this._colorFunction(), + } as unknown as $3Dmol.StickStyleSpec; } return style; } + /** + * Get the values that should be used to color atoms when coloring them + * according to properties + */ + private _colorValues(property: string, transform: string): Array { + assert(this._properties !== undefined); + assert(Object.keys(this._properties).includes(property)); + + // JSON.parse & JSON.stringify to make a deep copy of the properties to + // avoid modifying the original ones. + // + // This also transforms all `undefined` values into `null` + let currentProperty = JSON.parse(JSON.stringify(this._properties[property])) as Array< + number | null + >; + + let transformFunc = (x: number) => x; + if (transform === 'log') { + transformFunc = Math.log10; + } else if (transform === 'sqrt') { + transformFunc = Math.sqrt; + } else if (transform === 'inverse') { + transformFunc = (x) => 1 / x; + } + + currentProperty = currentProperty.map((value) => { + if (value !== null && !isNaN(value)) { + return transformFunc(value); + } else { + return value; + } + }); + + return currentProperty; + } + + /** + * Get a function computing the atom color, that can be used as 3Dmol + * `colorfunc` + */ + private _colorFunction(): ((atom: $3Dmol.AtomSpec) => number) | undefined { + if (this._properties === undefined) { + return undefined; + } + + const property = this._options.color.property.value; + if (property === 'element') { + return undefined; + } + + const transform = this._options.color.transform.value; + + const currentProperty = this._colorValues(property, transform); + const [min, max]: [number, number] = [ + this._options.color.min.value, + this._options.color.max.value, + ]; + + let grad: $3Dmol.Gradient = new $3Dmol.Gradient.RWB(max, min); + + if (this._options.color.palette.value === 'BWR') { + // min and max are swapped to ensure red is used for high values, + // blue for low values + grad = new $3Dmol.Gradient.RWB(max, min); + } else if (this._options.color.palette.value === 'ROYGB') { + grad = new $3Dmol.Gradient.ROYGB(min, max); + } else if (this._options.color.palette.value === 'Sinebow') { + grad = new $3Dmol.Gradient.Sinebow(min, max); + } + + return (atom: $3Dmol.AtomSpec) => { + assert(atom.serial !== undefined); + const value = currentProperty[atom.serial]; + if (value === null) { + // missing values + return 0xdddddd; + } else if (isNaN(value)) { + // NaN values + return 0x222222; + } else { + return grad.valueToHex(value); + } + }; + } + /** * Get the style specification for the hidden/background atoms when * highlighting a specific environment @@ -1069,7 +1293,7 @@ export class MoleculeViewer { if (bgStyle === 'ball-stick' && this._options.shape.value === '') { style.sphere = { // slightly smaller scale than the main style - scale: 0.219, + scale: this._options.spaceFilling.value ? 0.999 : 0.219, opacity: defaultOpacity(), }; } @@ -1088,6 +1312,25 @@ export class MoleculeViewer { if (style.sphere !== undefined) { style.sphere.color = 0x808080; } + } else if (bgColor === 'prop') { + if (style.stick !== undefined) { + style.stick = { + // slightly smaller radius than the main style + radius: 0.149, + opacity: defaultOpacity(), + hidden: !this._options.bonds.value, + colorfunc: this._colorFunction(), + } as unknown as $3Dmol.StickStyleSpec; + } + + if (style.sphere !== undefined) { + style.sphere = { + // slightly smaller scale than the main style + scale: this._options.spaceFilling.value ? 0.999 : 0.219, + opacity: defaultOpacity(), + colorfunc: this._colorFunction(), + } as unknown as $3Dmol.SphereStyleSpec; + } } else { unreachable(); } @@ -1310,4 +1553,16 @@ export class MoleculeViewer { }), }; } + + /** Changes the step of the arrow buttons in min/max input based on dataset range*/ + private _setScaleStep(axisBounds: number[]): void { + if (axisBounds !== undefined) { + // round to 10 decimal places so it does not break in Firefox + const step = Math.round(((axisBounds[1] - axisBounds[0]) / 20) * 10 ** 10) / 10 ** 10; + const minElement = this._options.getModalElement(`atom-color-min`); + const maxElement = this._options.getModalElement(`atom-color-max`); + minElement.step = `${step}`; + maxElement.step = `${step}`; + } + } } diff --git a/src/utils/index.ts b/src/utils/index.ts index 7e10f65a8..8a55964d6 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -112,10 +112,14 @@ export function getFirstKey(map: Map, excluding?: K): K { } // get the max/min of an array. Math.min(...array) fails with very large arrays -export function arrayMaxMin(values: number[]): { max: number; min: number } { +export function arrayMaxMin(values: Array): { max: number; min: number } { let max = Number.NEGATIVE_INFINITY; let min = Number.POSITIVE_INFINITY; for (const value of values) { + if (value === null) { + continue; + } + if (value > max && isFinite(value)) { max = value; }