from mmtf.codecs import encode_array import msgpack from mmtf.utils import constants def make_entity_dict(chain_indices,sequence,description,entity_type): out_d = {} out_d["description"] = description out_d["type"] = entity_type out_d["chainIndexList"] = chain_indices out_d["sequence"] = sequence return out_d class Group(object): def __eq__(self, other): """Function to define equality""" if self.atom_name_list != other.atom_name_list: return False if self.charge_list != other.charge_list: return False if self.element_list != other.element_list: return False if self.group_type != other.group_type: return False if self.group_name != other.group_name: return False if self.single_letter_code != other.single_letter_code: return False if self.bond_atom_list != other.bond_atom_list: return False if self.bond_order_list != other.bond_order_list: return False return True def __init__(self): self.atom_name_list = [] self.bond_order_list = [] self.bond_atom_list = [] self.charge_list = [] self.element_list = [] self.group_name = constants.UNKOWN_GROUP_NAME self.group_type = constants.UNKOWN_GROUP_TYPE self.single_letter_code = constants.UNKNOWN_SL def convert_to_dict(self): """Convert the group object to an appropriate DICT""" out_dict = {} out_dict["groupName"] = self.group_name out_dict["atomNameList"] = self.atom_name_list out_dict["elementList"] = self.element_list out_dict["bondOrderList"] = self.bond_order_list out_dict["bondAtomList"] = self.bond_atom_list out_dict["formalChargeList"] = self.charge_list out_dict["singleLetterCode"] = self.single_letter_code out_dict["chemCompType"] = self.group_type return out_dict def get_unique_groups(input_list): """Function to get a unique list of groups.""" out_list = [] for item in input_list: if item not in out_list: out_list.append(item) return out_list class TemplateEncoder(object): """Template class to be used by third parties to pass data into other data structures.""" def init_structure(self, total_num_bonds, total_num_atoms, total_num_groups, total_num_chains, total_num_models, structure_id): """Initialise the structure object. :param total_num_bonds: the number of bonds in the structure :param total_num_atoms: the number of atoms in the structure :param total_num_groups: the number of groups in the structure :param total_num_chains: the number of chains in the structure :param total_num_models: the number of models in the structure :param structure_id the: id of the structure (e.g. PDB id) """ raise NotImplementedError def set_atom_info(self, atom_name, serial_number, alternative_location_id, x, y, z, occupancy, temperature_factor, element, charge): """Create an atom object an set the information. :param atom_name: the atom name, e.g. CA for this atom :param serial_number: the serial id of the atom (e.g. 1) :param alternative_location_id: the alternative location id for the atom, if present :param x: the x coordiante of the atom :param y: the y coordinate of the atom :param z: the z coordinate of the atom :param occupancy: the occupancy of the atom :param temperature_factor: the temperature factor of the atom :param element: the element of the atom, e.g. C for carbon. According to IUPAC. Calcium is Ca :param charge: the formal atomic charge of the atom """ raise NotImplementedError def set_chain_info(self, chain_id, chain_name, num_groups): """Set the chain information. :param chain_id: the asym chain id from mmCIF :param chain_name: the auth chain id from mmCIF :param num_groups: the number of groups this chain has """ raise NotImplementedError def set_entity_info(self, chain_indices, sequence, description, entity_type): """Set the entity level information for the structure. :param chain_indices: the indices of the chains for this entity :param sequence: the one letter code sequence for this entity :param description: the description for this entity :param entity_type: the entity type (polymer,non-polymer,water) """ raise NotImplementedError def set_group_info(self, group_name, group_number, insertion_code, group_type, atom_count, bond_count, single_letter_code, sequence_index, secondary_structure_type): """Set the information for a group :param group_name: the name of this group,e.g. LYS :param group_number: the residue number of this group :param insertion_code: the insertion code for this group :param group_type: a string indicating the type of group (as found in the chemcomp dictionary. Empty string if none available. :param atom_count: the number of atoms in the group :param bond_count: the number of unique bonds in the group :param single_letter_code: the single letter code of the group :param sequence_index: the index of this group in the sequence defined by the enttiy :param secondary_structure_type: the type of secondary structure used (types are according to DSSP and number to type mappings are defined in the specification) """ raise NotImplementedError def set_model_info(self, model_id, chain_count): # FIXME model_id here is meaningles and potentially misleading. """Set the information for a model. :param model_id: the index for the model :param chain_count: the number of chains in the model """ raise NotImplementedError def set_xtal_info(self, space_group, unit_cell): """Set the crystallographic information for the structure :param space_group: the space group name, e.g. "P 21 21 21" :param unit_cell: an array of length 6 with the unit cell parameters in order: a, b, c, alpha, beta, gamma """ raise NotImplementedError def set_header_info(self, r_free, r_work, resolution, title, deposition_date, release_date, experimental_methods): """Sets the header information. :param r_free: the measured R-Free for the structure :param r_work: the measure R-Work for the structure :param resolution: the resolution of the structure :param title: the title of the structure :param deposition_date: the deposition date of the structure :param release_date: the release date of the structure :param experimnetal_methods: the list of experimental methods in the structure """ raise NotImplementedError def set_bio_assembly_trans(self, bio_assembly_index, input_chain_indices, input_transform): """Set the Bioassembly transformation information. A single bioassembly can have multiple transforms, :param bio_assembly_index: the integer index of the bioassembly :param input_chain_indices: the list of integer indices for the chains of this bioassembly :param input_transformation: the list of doubles for the transform of this bioassmbly transform""" raise NotImplementedError def finalize_structure(self): """Any functions needed to cleanup the structure.""" raise NotImplementedError def set_group_bond(self, atom_index_one, atom_index_two, bond_order): """Add bonds within a group. :param atom_index_one: the integer atom index (in the group) of the first partner in the bond :param atom_index_two: the integer atom index (in the group) of the second partner in the bond :param bond_order: the integer bond order """ raise NotImplementedError def set_inter_group_bond(self, atom_index_one, atom_index_two, bond_order): """Add bonds between groups. :param atom_index_one: the integer atom index (in the structure) of the first partner in the bond :param atom_index_two: the integer atom index (in the structure) of the second partner in the bond :param bond_order the bond order """ raise NotImplementedError class MMTFEncoder(TemplateEncoder): def encode_data(self): """Encode the data back into a dict.""" output_data = {} output_data["groupTypeList"] = encode_array(self.group_type_list, 4, 0) output_data["xCoordList"] = encode_array(self.x_coord_list, 10, 1000) output_data["yCoordList"] = encode_array(self.y_coord_list, 10, 1000) output_data["zCoordList"] = encode_array(self.z_coord_list, 10, 1000) output_data["bFactorList"] = encode_array(self.b_factor_list, 10, 100) output_data["occupancyList"] = encode_array(self.occupancy_list, 9, 100) output_data["atomIdList"] = encode_array(self.atom_id_list, 8, 0) output_data["altLocList"] = encode_array(self.alt_loc_list, 6, 0) output_data["insCodeList"] = encode_array(self.ins_code_list, 6, 0) output_data["groupIdList"] = encode_array(self.group_id_list, 8, 0) output_data["groupList"] = self.group_list output_data["sequenceIndexList"] = encode_array(self.sequence_index_list, 8, 0) output_data["chainNameList"] = encode_array(self.chain_name_list, 5, 4) output_data["chainIdList"] = encode_array(self.chain_id_list, 5, 4) output_data["bondAtomList"] = encode_array(self.bond_atom_list, 4, 0) output_data["bondOrderList"] = encode_array(self.bond_order_list, 2, 0) output_data["secStructList"] = encode_array(self.sec_struct_list, 2, 0) output_data["chainsPerModel"] = self.chains_per_model output_data["groupsPerChain"] = self.groups_per_chain output_data["spaceGroup"] = self.space_group output_data["mmtfVersion"] = self.mmtf_version output_data["mmtfProducer"] = self.mmtf_producer output_data["structureId"] = self.structure_id output_data["entityList"] = self.entity_list output_data["bioAssemblyList"] = self.bio_assembly output_data["rFree"] = self.r_free output_data["rWork"] = self.r_work output_data["resolution"] = self.resolution output_data["title"] = self.title output_data["experimentalMethods"] = self.experimental_methods output_data["depositionDate"] = self.deposition_date output_data["releaseDate"] = self.release_date output_data["unitCell"] = self.unit_cell output_data["numBonds"] = self.num_bonds output_data["numChains"] = self.num_chains output_data["numModels"] = self.num_models output_data["numAtoms"] = self.num_atoms output_data["numGroups"] = self.num_groups return output_data def get_msgpack(self): """Get the msgpack of the encoded data.""" return msgpack.packb(self.encode_data(), use_bin_type=True) def write_file(self, file_path): with open(file_path, "wb") as out_f: out_f.write(self.get_msgpack()) def init_structure(self, total_num_bonds, total_num_atoms, total_num_groups, total_num_chains, total_num_models, structure_id): """Initialise the structure object. :param total_num_bonds: the number of bonds in the structure :param total_num_atoms: the number of atoms in the structure :param total_num_groups: the number of groups in the structure :param total_num_chains: the number of chains in the structure :param total_num_models: the number of models in the structure :param structure_id the: id of the structure (e.g. PDB id) """ self.mmtf_version = constants.MMTF_VERSION self.mmtf_producer = constants.PRODUCER self.num_atoms = total_num_atoms self.num_bonds = total_num_bonds self.num_groups = total_num_groups self.num_chains = total_num_chains self.num_models = total_num_models self.structure_id = structure_id # initialise the arrays self.x_coord_list = [] self.y_coord_list = [] self.z_coord_list = [] self.group_type_list = [] self.entity_list = [] self.b_factor_list = [] self.occupancy_list = [] self.atom_id_list = [] self.alt_loc_list = [] self.ins_code_list = [] self.group_id_list = [] self.sequence_index_list = [] self.group_list = [] self.chain_name_list = [] self.chain_id_list = [] self.bond_atom_list = [] self.bond_order_list = [] self.sec_struct_list = [] self.chains_per_model = [] self.groups_per_chain = [] self.current_group = None self.bio_assembly = [] def set_atom_info(self, atom_name, serial_number, alternative_location_id, x, y, z, occupancy, temperature_factor, element, charge): """Create an atom object an set the information. :param atom_name: the atom name, e.g. CA for this atom :param serial_number: the serial id of the atom (e.g. 1) :param alternative_location_id: the alternative location id for the atom, if present :param x: the x coordiante of the atom :param y: the y coordinate of the atom :param z: the z coordinate of the atom :param occupancy: the occupancy of the atom :param temperature_factor: the temperature factor of the atom :param element: the element of the atom, e.g. C for carbon. According to IUPAC. Calcium is Ca :param charge: the formal atomic charge of the atom """ self.x_coord_list.append(x) self.y_coord_list.append(y) self.z_coord_list.append(z) self.atom_id_list.append(serial_number) self.alt_loc_list.append(alternative_location_id) self.occupancy_list.append(occupancy) self.b_factor_list.append(temperature_factor) ## Now add the group level data self.current_group.atom_name_list.append(atom_name) self.current_group.charge_list.append(charge) self.current_group.element_list.append(element) def set_chain_info(self, chain_id, chain_name, num_groups): """Set the chain information. :param chain_id: the asym chain id from mmCIF :param chain_name: the auth chain id from mmCIF :param num_groups: the number of groups this chain has """ self.chain_id_list.append(chain_id) self.chain_name_list.append(chain_name) self.groups_per_chain.append(num_groups) def set_entity_info(self, chain_indices, sequence, description, entity_type): """Set the entity level information for the structure. :param chain_indices: the indices of the chains for this entity :param sequence: the one letter code sequence for this entity :param description: the description for this entity :param entity_type: the entity type (polymer,non-polymer,water) """ self.entity_list.append(make_entity_dict(chain_indices,sequence,description,entity_type)) def set_group_info(self, group_name, group_number, insertion_code, group_type, atom_count, bond_count, single_letter_code, sequence_index, secondary_structure_type): """Set the information for a group :param group_name: the name of this group,e.g. LYS :param group_number: the residue number of this group :param insertion_code: the insertion code for this group :param group_type: a string indicating the type of group (as found in the chemcomp dictionary. Empty string if none available. :param atom_count: the number of atoms in the group :param bond_count: the number of unique bonds in the group :param single_letter_code: the single letter code of the group :param sequence_index: the index of this group in the sequence defined by the enttiy :param secondary_structure_type: the type of secondary structure used (types are according to DSSP and number to type mappings are defined in the specification) """ # Add the group to the overall list - unless it's the first time round if self.current_group is not None: self.group_list.append(self.current_group) # Add the group level information self.group_id_list.append(group_number) self.ins_code_list.append(insertion_code) self.sequence_index_list.append(sequence_index) self.sec_struct_list.append(secondary_structure_type) self.current_group = Group() self.current_group.group_name = group_name self.current_group.group_type = group_type self.current_group.single_letter_code = single_letter_code def set_model_info(self, model_id, chain_count): # FIXME model_id here is meaningles and potentially misleading. """Set the information for a model. :param model_id: the index for the model :param chain_count: the number of chains in the model """ self.chains_per_model.append(chain_count) def set_xtal_info(self, space_group, unit_cell): """Set the crystallographic information for the structure :param space_group: the space group name, e.g. "P 21 21 21" :param unit_cell: an array of length 6 with the unit cell parameters in order: a, b, c, alpha, beta, gamma """ self.space_group = space_group self.unit_cell = unit_cell def set_header_info(self, r_free, r_work, resolution, title, deposition_date, release_date, experimental_methods): """Sets the header information. :param r_free: the measured R-Free for the structure :param r_work: the measure R-Work for the structure :param resolution: the resolution of the structure :param title: the title of the structure :param deposition_date: the deposition date of the structure :param release_date: the release date of the structure :param experimnetal_methods: the list of experimental methods in the structure """ self.r_free = r_free self.r_work = r_work self.resolution = resolution self.title = title self.deposition_date = deposition_date self.release_date = release_date self.experimental_methods = experimental_methods def set_bio_assembly_trans(self, bio_assembly_index, input_chain_indices, input_transform): """Set the Bioassembly transformation information. A single bioassembly can have multiple transforms, :param bio_assembly_index: the integer index of the bioassembly :param input_chain_indices: the list of integer indices for the chains of this bioassembly :param input_transformation: the list of doubles for the transform of this bioassmbly transform""" this_bioass = None for bioass in self.bio_assembly: if bioass['name'] == str(bio_assembly_index): this_bioass = bioass break if not this_bioass: this_bioass = {"name": str(bio_assembly_index), 'transformList': []} else: self.bio_assembly.remove(this_bioass) this_bioass['transformList'].append({'chainIndexList':input_chain_indices,'matrix': input_transform}) self.bio_assembly.append(this_bioass) def finalize_structure(self): """Any functions needed to cleanup the structure.""" self.group_list.append(self.current_group) group_set = get_unique_groups(self.group_list) for item in self.group_list: self.group_type_list.append(group_set.index(item)) self.group_list = [x.convert_to_dict() for x in group_set] def set_group_bond(self, atom_index_one, atom_index_two, bond_order): """Add bonds within a group. :param atom_index_one: the integer atom index (in the group) of the first partner in the bond :param atom_index_two: the integer atom index (in the group) of the second partner in the bond :param bond_order: the integer bond order """ self.current_group.bond_atom_list.append(atom_index_one) self.current_group.bond_atom_list.append(atom_index_two) self.current_group.bond_order_list.append(bond_order) def set_inter_group_bond(self, atom_index_one, atom_index_two, bond_order): """Add bonds between groups. :param atom_index_one: the integer atom index (in the structure) of the first partner in the bond :param atom_index_two: the integer atom index (in the structure) of the second partner in the bond :param bond_order the bond order """ self.bond_atom_list.append(atom_index_one) self.bond_atom_list.append(atom_index_two) self.bond_order_list.append(bond_order)