mol-model: filter chem_comp in CIF export

src/mol-model/structure/export/categories/misc.ts

+/**
+ * Copyright (c) 2017-2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { Column } from 'mol-data/db';
+import { mmCIF_Database, mmCIF_Schema } from 'mol-io/reader/cif/schema/mmcif';
+import { CifWriter } from 'mol-io/writer/cif';
+import { unionMany } from 'mol-util/set';
+import { Model } from '../../model';
+import { CifExportContext } from '../mmcif';
+import CifCategory = CifWriter.Category
+import { Structure } from '../../structure';
+
+export const _chem_comp: CifCategory<CifExportContext> = {
+    name: 'chem_comp',
+    instance({ structures, cache }) {
+        const chem_comp = getCifCategory(structures[0].model, 'chem_comp');
+        if (!chem_comp) return CifCategory.Empty;
+        const { id } = chem_comp;
+        const names = cache.uniqueResidueNames || (cache.uniqueResidueNames = getUniqueResidueNames(structures));
+        const indices = Column.indicesOf(id, id => names.has(id));
+        return CifCategory.ofTable(chem_comp, indices);
+    }
+}
+
+export const _pdbx_chem_comp_identifier: CifCategory<CifExportContext> = {
+    name: 'pdbx_chem_comp_identifier',
+    instance({ structures, cache }) {
+        const pdbx_chem_comp_identifier = getCifCategory(structures[0].model, 'pdbx_chem_comp_identifier');
+        if (!pdbx_chem_comp_identifier) return CifCategory.Empty;
+        const { comp_id } = pdbx_chem_comp_identifier;
+        const names = cache.uniqueResidueNames || (cache.uniqueResidueNames = getUniqueResidueNames(structures));
+        const indices = Column.indicesOf(comp_id, id => names.has(id));
+        return CifCategory.ofTable(pdbx_chem_comp_identifier, indices);
+    }
+}
+
+function getCifCategory<K extends keyof mmCIF_Schema>(model: Model, name: K): mmCIF_Database[K] | undefined {
+    if (model.sourceData.kind !== 'mmCIF') return;
+    return model.sourceData.data[name];
+}
+
+function getUniqueResidueNames(structures: Structure[]) {
+    return unionMany(structures.map(s => s.uniqueResidueNames));
+}
\ No newline at end of file

src/mol-model/structure/export/mmcif.ts

@@ -12,16 +12,21 @@ import { _atom_site } from './categories/atom_site';
 import CifCategory = CifWriter.Category
 import { _struct_conf, _struct_sheet_range } from './categories/secondary-structure';
 import { _pdbx_struct_mod_residue } from './categories/modified-residues';
+import { _chem_comp, _pdbx_chem_comp_identifier } from './categories/misc';
+import { Model } from '../model';
 
 export interface CifExportContext {
     structures: Structure[],
+    firstModel: Model,
     cache: any
 }
 
 export namespace CifExportContext {
     export function create(structures: Structure | Structure[]): CifExportContext {
+        const structureArray = Array.isArray(structures) ? structures : [structures];
         return {
-            structures: Array.isArray(structures) ? structures : [structures],
+            structures: structureArray,
+            firstModel: structureArray[0].model,
             cache: Object.create(null)
         };
     }
@@ -68,9 +73,9 @@ const Categories = [
     _struct_sheet_range,
 
     // Sequence
-    copy_mmCif_category('struct_asym'), // TODO: filter only present chains?
-    copy_mmCif_category('entity_poly'),
-    copy_mmCif_category('entity_poly_seq'),
+    copy_mmCif_category('struct_asym'), // TODO: filter only present entities?
+    copy_mmCif_category('entity_poly'), // TODO: filter only present entities?
+    copy_mmCif_category('entity_poly_seq'), // TODO: filter only present entities?
 
     // Branch
     copy_mmCif_category('pdbx_entity_branch'),
@@ -79,8 +84,8 @@ const Categories = [
 
     // Misc
     // TODO: filter for actual present residues?
-    copy_mmCif_category('chem_comp'),
-    copy_mmCif_category('pdbx_chem_comp_identifier'),
+    _chem_comp,
+    _pdbx_chem_comp_identifier,
     copy_mmCif_category('atom_sites'),
 
     _pdbx_struct_mod_residue,

src/mol-model/structure/structure/structure.ts

@@ -133,7 +133,7 @@ class Structure {
         return this._props.models;
     }
 
-    get uniqueResidueName() {
+    get uniqueResidueNames() {
         return this._props.uniqueResidueNames
             || (this._props.uniqueResidueNames = getUniqueResidueNames(this));
     }

src/mol-util/set.ts

@@ -21,6 +21,16 @@ export function union<T>(setA: Set<T>, setB: Set<T>): Set<T> {
     return union;
 }
 
+export function unionMany<T>(sets: Set<T>[]) {
+    if (sets.length === 0) return new Set<T>();
+    if (sets.length === 1) return sets[0];
+    const union = new Set(sets[0]);
+    for (let i = 1; i < sets.length; i++) {
+        for (const elem of Array.from(sets[i])) union.add(elem);
+    }
+    return union;
+}
+
 /** Create set containing elements of set a that are also in set b. */
 export function intersection<T>(setA: Set<T>, setB: Set<T>): Set<T> {
     const intersection = new Set();