"""Functions to create PDB files from entities optionally with data in the B-factor and Q (occupancy) columns. """ from collections import defaultdict from itertools import chain from cogent.struct.selection import einput from cogent.data.protein_properties import AA_NAMES from cogent.parse.pdb import dict2pdb, dict2ter __author__ = "Marcin Cieslik" __copyright__ = "Copyright 2007-2012, The Cogent Project" __credits__ = ["Marcin Cieslik"] __license__ = "GPL" __version__ = "1.5.3-dev" __maintainer__ = "Marcin Cieslik" __email__ = "mpc4p@virginia.edu" __status__ = "Development" def write_header(header): if not isinstance(header, dict): return (header or []) xt = ('REMARK 200 EXPERIMENT TYPE : %s\n', 'experiment_type') sg = ('REMARK 290 SYMMETRY OPERATORS FOR SPACE GROUP: %s\n', 'space_group') tests = [xt, sg] results = [] for test in tests: string, key = test try: value = header[key] results.append(string % value) except (KeyError, TypeError): pass return results def write_coords(atoms): old_fields = defaultdict(str) coords = ['MODEL 1\n'] for atom in atoms.sortedvalues(): fields = atom.getDict() if (old_fields['chain_id'] != fields['chain_id']) and old_fields['chain_id'] and old_fields['res_name'] in AA_NAMES: coords.append(dict2ter(old_fields)) if (old_fields['model'] != fields['model']) and old_fields['model'] != '': # model can be 0 :) if old_fields['chain_id'] and old_fields['res_name'] in AA_NAMES: coords.append(dict2ter(old_fields)) coords.append('ENDMDL\n') coords.append('MODEL %4i\n' % (fields['model'] + 1)) coords.append(dict2pdb(fields)) old_fields = fields if fields['res_name'] in AA_NAMES: coords.append(dict2ter(fields)) coords.append('ENDMDL\n') coords.append('END \n') return coords def write_trailer(trailer): if not isinstance(trailer, dict): return (trailer or []) def number(data, rest_val): try: return float(data) except TypeError: return rest_val def iterable(data, rest_val): try: return float(len(data)) except TypeError: return rest_val def PDBWriter(f, entities, header_=None, trailer_=None): structure = einput(entities, level='A', name='atoms') # hierarchy: args, dicts try: header = (header_ or entities.raw_header or entities.header) except AttributeError: header = header_ try: trailer = (trailer_ or entities.raw_trailer or entities.trailer) except AttributeError: trailer = trailer_ coords = write_coords(structure) header = write_header(header) trailer = write_trailer(trailer) for part in chain(header, coords, trailer): f.writelines(part) # did not open do not close # f.close() def PDBXWriter(f, entities, level, b_key, b_mode=None, b_val=0.0, q_key=None, \ q_mode=None, q_val=0.0): """Writes data from the ``xtra`` dictionary into B and Q columns of a PDB file. The level from which the dictionary is taken can be specified. The b_key and q_key specifies should be a key in the dictionaries, b_val and q_val are the default values, b_mode and q_mode can be "number" - ``float`` will be called to transform the data or "iterable" which will return the length (``len``) of the sequence. The B and Q columns can contain only numeric values, thus any data which we wish to store in those columns needs to be converted. The following functions convert data to numeric form. boolean type can also be treated as number. """ b_mode = (b_mode or 'number') # B q_mode = (q_mode or 'number') # Q entities = einput(entities, level) for entity in entities: q_data = eval(q_mode)(entity.xtra.get(q_key), b_val) # occupancy b_data = eval(b_mode)(entity.xtra.get(b_key), q_val) # b-factor if level != 'A': atoms = einput(entity, 'A') for atom in atoms: if b_key: atom.setBfactor(b_data) if q_key: atom.setOccupancy(q_data) else: if b_key: entity.setBfactor(b_data) if q_key: entity.setOccupancy(q_data) PDBWriter(f, entities) # we try to preserve the headers