introduce map, ditch getLabel

src/mol-io/writer/ligand-encoder.ts

@@ -86,14 +86,15 @@ export abstract class LigandEncoder implements Encoder<string> {
         return StringBuilder.getString(this.builder);
     }
 
-    protected getAtoms<Ctx>(instance: Category.Instance<Ctx>, source: any): Atom[] {
-        const sortedFields = this.getSortedFields(instance, ['Cartn_x', 'Cartn_y', 'Cartn_z', 'type_symbol']);
+    protected getAtoms<Ctx>(instance: Category.Instance<Ctx>, source: any): Map<string, Atom> {
+        const sortedFields = this.getSortedFields(instance, ['Cartn_x', 'Cartn_y', 'Cartn_z']);
         const label_atom_id = this.getField(instance, 'label_atom_id');
-        return this._getAtoms(source, sortedFields, label_atom_id);
+        const type_symbol = this.getField(instance, 'type_symbol');
+        return this._getAtoms(source, sortedFields, label_atom_id, type_symbol);
     }
 
-    private _getAtoms(source: any, fields: Field<any, any>[], label_atom_id: Field<any, any>): Atom[] {
-        const atoms: Atom[] = [];
+    private _getAtoms(source: any, fields: Field<any, any>[], label_atom_id: Field<any, any>, type_symbol: Field<any, any>): Map<string, Atom> {
+        const atoms = new Map<string, Atom>();
         let index = 0;
 
         // is outer loop even needed?
@@ -108,8 +109,8 @@ export abstract class LigandEncoder implements Encoder<string> {
                 const key = it.move();
 
                 const lai = label_atom_id.value(key, data, index) as string;
-                const label = this.getLabel(lai);
-                if (this.skipHydrogen(label)) {
+                const ts = type_symbol.value(key, data, index) as string;
+                if (this.skipHydrogen(ts)) {
                     index++;
                     continue;
                 }
@@ -119,8 +120,9 @@ export abstract class LigandEncoder implements Encoder<string> {
                     const f: Field<any, any> = fields[_f]!;
                     a[f.name] = f.value(key, data, index);
                 }
+                a[type_symbol.name] = ts;
 
-                atoms.push(Atom(a));
+                atoms.set(lai, Atom(a));
                 index++;
             }
         }
@@ -128,16 +130,11 @@ export abstract class LigandEncoder implements Encoder<string> {
         return atoms;
     }
 
-    protected skipHydrogen(label: string) {
+    protected skipHydrogen(type_symbol: string) {
         if (this.hydrogens) {
             return false;
         }
-        return label.startsWith('H');
-    }
-
-    protected getLabel(s: string) {
-        // actually, getTypeSymbol would be way more useful
-        return s.replace(/[^A-Z]+/g, '');
+        return type_symbol === 'H';
     }
 
     private getSortedFields<Ctx>(instance: Category.Instance<Ctx>, names: string[]) {