support atom id list in selection helper

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

@@ -13,7 +13,7 @@ import { StructureComponentManager } from '../../mol-plugin-state/manager/struct
 import { StructureComponentRef, StructureRef } from '../../mol-plugin-state/manager/structure/hierarchy-state';
 import { StructureSelectionModifier } from '../../mol-plugin-state/manager/structure/selection';
 import { PluginContext } from '../../mol-plugin/context';
-import { compileResidueListSelection } from '../../mol-script/util/residue-list';
+import { compileIdListSelection } from '../../mol-script/util/id-list';
 import { memoizeLatest } from '../../mol-util/memoize';
 import { ParamDefinition } from '../../mol-util/param-definition';
 import { capitalize, stripTags } from '../../mol-util/string';
@@ -172,7 +172,7 @@ export class StructureSelectionActionsControls extends PluginUIComponent<{}, Str
         //       the plan is to add support to input queries in different languages
         //       after this has been implemented in mol-script
         const helpers = [
-            { kind: 'residue-list' as SelectionHelperType, category: 'Helpers', label: 'Residue List', description: 'Create a selection from a list of residue ranges.' }
+            { kind: 'residue-list' as SelectionHelperType, category: 'Helpers', label: 'Atom/Residue Identifier List', description: 'Create a selection from a list of atom/residue ranges.' }
         ];
         this.helpersItems = ActionMenu.createItems(helpers, {
             label: q => q.label,
@@ -252,7 +252,7 @@ export class StructureSelectionActionsControls extends PluginUIComponent<{}, Str
         } else if (ActionHeader.has(this.state.action as any) && this.state.helper === 'residue-list') {
             const close = () => this.setState({ action: void 0, helper: void 0 });
             children = <div className='msp-selection-viewport-controls-actions'>
-                <ControlGroup header='Residue List' title='Click to close.' initialExpanded={true} hideExpander={true} hideOffset={true} onHeaderClick={close} topRightIcon={CloseSvg}>
+                <ControlGroup header='Atom/Residue Identifier List' title='Click to close.' initialExpanded={true} hideExpander={true} hideOffset={true} onHeaderClick={close} topRightIcon={CloseSvg}>
                     <ResidueListSelectionHelper modifier={this.state.action as any} plugin={this.plugin} close={close} />
                 </ControlGroup>
             </div>;
@@ -384,8 +384,8 @@ class ApplyThemeControls extends PurePluginUIComponent<ApplyThemeControlsProps,
 }
 
 const ResidueListIdTypeParams = {
-    idType: ParamDefinition.Select<'auth' | 'label'>('auth', ParamDefinition.arrayToOptions(['auth', 'label'])),
-    residues: ParamDefinition.Text('', { description: 'A comma separated list of residue ranges in given chain, e.g. A 10-15, B 25, C 30:i' })
+    idType: ParamDefinition.Select<'auth' | 'label' | 'atom-id'>('auth', ParamDefinition.arrayToOptions(['auth', 'label', 'atom-id'])),
+    identifiers: ParamDefinition.Text('', { description: 'A comma separated list of atom identifiers (e.g. 10, 15-25) or residue ranges in given chain (e.g. A 10-15, B 25, C 30:i)' })
 };
 
 const DefaultResidueListIdTypeParams = ParamDefinition.getDefaultValues(ResidueListIdTypeParams);
@@ -394,11 +394,11 @@ function ResidueListSelectionHelper({ modifier, plugin, close }: { modifier: Str
     const [state, setState] = React.useState(DefaultResidueListIdTypeParams);
 
     const apply = () => {
-        if (state.residues.length === 0) return;
+        if (state.identifiers.trim().length === 0) return;
 
         try {
             close();
-            const query = compileResidueListSelection(state.residues, state.idType);
+            const query = compileIdListSelection(state.identifiers, state.idType);
             plugin.managers.structure.selection.fromCompiledQuery(modifier, query, false);
         } catch (e) {
             plugin.log.error(`Failed to create selection: ${e}`);
@@ -407,7 +407,7 @@ function ResidueListSelectionHelper({ modifier, plugin, close }: { modifier: Str
 
     return <>
         <ParameterControls params={ResidueListIdTypeParams} values={state} onChangeValues={setState} onEnter={apply} />
-        <Button className='msp-btn-commit msp-btn-commit-on' disabled={state.residues.length === 0} onClick={apply} style={{ marginTop: '1px' }}>
+        <Button className='msp-btn-commit msp-btn-commit-on' disabled={state.identifiers.trim().length === 0} onClick={apply} style={{ marginTop: '1px' }}>
             {capitalize(modifier)} Selection
         </Button>
     </>;

src/mol-script/util/residue-list.ts → src/mol-script/util/id-list.ts

@@ -8,6 +8,7 @@ import { StructureQuery } from '../../mol-model/structure/query';
 import { Expression } from '../language/expression';
 import { MolScriptBuilder as MS } from '../language/builder';
 import { compile } from '../runtime/query/base';
+import { UniqueArray } from '../../mol-data/generic';
 
 // TODO: make this into a separate "language"?
 
@@ -15,7 +16,7 @@ type ResidueListSelectionEntry =
     | { kind: 'single', asym_id: string; seq_id: number; ins_code?: string }
     | { kind: 'range', asym_id: string; seq_id_beg: number; seq_id_end: number; }
 
-function entriesToQuery(xs: ResidueListSelectionEntry[], kind: 'auth' | 'label') {
+function residueEntriesToQuery(xs: ResidueListSelectionEntry[], kind: 'auth' | 'label') {
     const groups: Expression[] = [];
 
     const asym_id_key = kind === 'auth' ? 'auth_asym_id' as const : 'label_asym_id' as const;
@@ -45,6 +46,23 @@ function entriesToQuery(xs: ResidueListSelectionEntry[], kind: 'auth' | 'label')
     return compile(query) as StructureQuery;
 }
 
+function atomEntriesToQuery(xs: [number, number][]) {
+    const set = UniqueArray.create<number>();
+
+    for (const [a, b] of xs) {
+        for (let i = a; i <= b; i++) {
+            UniqueArray.add(set, i, i);
+        }
+    }
+
+    const query = MS.struct.generator.atomGroups({
+        'atom-test': MS.core.set.has([MS.set(...set.array), MS.ammp('id')])
+    });
+
+    return compile(query) as StructureQuery;
+}
+
+
 function parseRange(c: string, s: string[], e: number): ResidueListSelectionEntry | undefined {
     if (!c || s.length === 0 || Number.isNaN(+s[0])) return;
     if (Number.isNaN(e)) {
@@ -65,8 +83,20 @@ function parseResidueListSelection(input: string): ResidueListSelectionEntry[] {
         .filter(e => !!e) as ResidueListSelectionEntry[];
 }
 
+function parseAtomListSelection(input: string): [number, number][] {
+    return input.split(',') // 1-3, 3 => [1-3, 3]
+        .map(e => e.trim().split(/\s+|[-]/g).filter(e => !!e)) // [1-3, 3] => [[1, 3], [3]]
+        .filter(e => e.length === 1 || e.length === 2)
+        .map(e => e.length === 1 ? [+e[0], +e[0]] : [+e[0], +e[1]]) as [number, number][];
+}
+
 // parses a list of residue ranges, e.g. A 10-100, B 30, C 12:i
-export function compileResidueListSelection(input: string, idType: 'auth' | 'label') {
-    const entries = parseResidueListSelection(input);
-    return entriesToQuery(entries, idType);
+export function compileIdListSelection(input: string, idType: 'auth' | 'label' | 'atom-id') {
+    if (idType === 'atom-id') {
+        const entries = parseAtomListSelection(input);
+        return atomEntriesToQuery(entries);
+    } else {
+        const entries = parseResidueListSelection(input);
+        return residueEntriesToQuery(entries, idType);
+    }
 }