// Copyright 2024 DeepMind Technologies Limited // // AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of // this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/ // // To request access to the AlphaFold 3 model parameters, follow the process set // out at https://github.com/google-deepmind/alphafold3. You may only use these // if received directly from Google. Use is subject to terms of use available at // https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md #include "alphafold3/parsers/cpp/cif_dict_lib.h" #include #include #include #include #include #include #include #include #include #include "absl/algorithm/container.h" #include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/container/node_hash_map.h" #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" #include "absl/strings/strip.h" namespace alphafold3 { namespace { bool IsQuote(const char symbol) { return symbol == '\'' || symbol == '"'; } bool IsWhitespace(const char symbol) { return symbol == ' ' || symbol == '\t'; } // Splits line into tokens, returns whether successful. bool SplitLineInline(absl::string_view line, std::vector* tokens) { // See https://www.iucr.org/resources/cif/spec/version1.1/cifsyntax for (int i = 0, line_length = line.length(); i < line_length;) { // Skip whitespace (spaces or tabs). while (IsWhitespace(line[i])) { if (++i == line_length) { break; } } if (i == line_length) { break; } // Skip comments (from # until the end of the line). If # is a non-comment // character, it must be inside a quoted token. if (line[i] == '#') { break; } int start_index; int end_index; if (IsQuote(line[i])) { // Token in single or double quotes. CIF v1.1 specification considers a // quote to be an opening quote only if it is at the beginning of a token. // So e.g. A' B has tokens A' and B. Also, ""A" is a token "A. const char quote_char = line[i++]; start_index = i; // Find matching quote. The double loop is not strictly necessary, but // optimises a bit better. while (true) { while (i < line_length && line[i] != quote_char) { ++i; } if (i == line_length) { // Reached the end of the line while still being inside a token. return false; } if (i + 1 == line_length || IsWhitespace(line[i + 1])) { break; } ++i; } end_index = i++; } else { // Non-quoted token. Read until reaching whitespace. start_index = i++; while (i < line_length && !IsWhitespace(line[i])) { ++i; } end_index = i; } tokens->push_back(line.substr(start_index, end_index - start_index)); } return true; } using HeapStrings = std::vector>; // The majority of strings can be viewed on original cif_string. // heap_strings store multi-line tokens that have internal white-space stripped. absl::StatusOr> TokenizeInternal( absl::string_view cif_string, HeapStrings* heap_strings) { const std::vector lines = absl::StrSplit(cif_string, '\n'); std::vector tokens; // Heuristic: Most lines in an mmCIF are _atom_site lines with 21 tokens. tokens.reserve(lines.size() * 21); int line_num = 0; while (line_num < lines.size()) { auto line = lines[line_num]; line_num++; if (line.empty() || line[0] == '#') { // Skip empty lines or lines that contain only comments. continue; } else if (line[0] == ';') { // Leading whitespace on each line must be preserved while trailing // whitespace may be stripped. std::vector multiline_tokens; // Strip the leading ";". multiline_tokens.push_back( absl::StripTrailingAsciiWhitespace(line.substr(1))); while (line_num < lines.size()) { auto multiline = absl::StripTrailingAsciiWhitespace(lines[line_num]); line_num++; if (!multiline.empty() && multiline[0] == ';') { break; } else if (line_num == lines.size()) { return absl::InvalidArgumentError( "Last multiline token is not terminated by a semicolon."); } multiline_tokens.push_back(multiline); } heap_strings->push_back( std::make_unique(absl::StrJoin(multiline_tokens, "\n"))); tokens.emplace_back(*heap_strings->back()); } else { if (!SplitLineInline(line, &tokens)) { return absl::InvalidArgumentError( absl::StrCat("Line ended with quote open: ", line)); } } } return tokens; } // Returns whether the token doesn't need any quoting. This is true if the token // isn't empty and contains only safe characters [A-Za-z0-9.?-]. bool IsTrivialToken(const absl::string_view value) { if (value.empty()) { return false; } return std::all_of(value.begin(), value.end(), [](char c) { return absl::ascii_isalnum(c) || c == '.' || c == '?' || c == '-'; }); } // Returns whether the token needs to be a multiline token. This happens if it // has a newline or both single and double quotes. bool IsMultiLineToken(const absl::string_view value) { bool has_single_quotes = false; bool has_double_quotes = false; for (const char c : value) { if (c == '\n') { return true; } else if (c == '\'') { has_single_quotes = true; } else if (c == '"') { has_double_quotes = true; } } return has_single_quotes && has_double_quotes; } absl::string_view GetEscapeQuote(const absl::string_view value) { // Empty values should not happen, but if so, they should be quoted. if (value.empty()) { return "\""; } // // Shortcut for the most common cases where no quoting needed. // if (std::all_of(value.begin(), value.end(), [](char c) { // return absl::ascii_isalnum(c) || c == '.' || c == '?' || c == '-'; // })) { // return ""; // } // The value must not start with one of these CIF keywords. if (absl::StartsWithIgnoreCase(value, "data_") || absl::StartsWithIgnoreCase(value, "loop_") || absl::StartsWithIgnoreCase(value, "save_") || absl::StartsWithIgnoreCase(value, "stop_") || absl::StartsWithIgnoreCase(value, "global_")) { return "\""; } // The first character must not be a special character. const char first = value.front(); if (first == '_' || first == '#' || first == '$' || first == '[' || first == ']' || first == ';') { return "\""; } // No quotes or whitespace allowed inside. Rare case when both double and // single quotes are present is handled by IsMultiLineToken. bool use_double_quote = true; bool use_single_quote = true; bool needs_quote = false; for (const char c : value) { if (c == ' ' || c == '\t') { needs_quote = true; } else if (c == '"') { needs_quote = true; use_double_quote = false; } else if (c == '\'') { needs_quote = true; use_single_quote = false; } } if (needs_quote && use_double_quote) { return "\""; } else if (needs_quote && use_single_quote) { return "'"; } return ""; } int RecordIndex(absl::string_view record) { if (record == "_entry") { return 0; // _entry is always first. } if (record == "_atom_site") { return 2; // _atom_site is always last. } return 1; // other records are between _entry and _atom_site. } struct RecordOrder { using is_transparent = void; // Enable heterogeneous lookup. bool operator()(absl::string_view lhs, absl::string_view rhs) const { std::size_t lhs_index = RecordIndex(lhs); std::size_t rhs_index = RecordIndex(rhs); return std::tie(lhs_index, lhs) < std::tie(rhs_index, rhs); } }; // Make sure the _atom_site loop columns are sorted in the PDB-standard way. constexpr absl::string_view kAtomSiteSortOrder[] = { "_atom_site.group_PDB", "_atom_site.id", "_atom_site.type_symbol", "_atom_site.label_atom_id", "_atom_site.label_alt_id", "_atom_site.label_comp_id", "_atom_site.label_asym_id", "_atom_site.label_entity_id", "_atom_site.label_seq_id", "_atom_site.pdbx_PDB_ins_code", "_atom_site.Cartn_x", "_atom_site.Cartn_y", "_atom_site.Cartn_z", "_atom_site.occupancy", "_atom_site.B_iso_or_equiv", "_atom_site.pdbx_formal_charge", "_atom_site.auth_seq_id", "_atom_site.auth_comp_id", "_atom_site.auth_asym_id", "_atom_site.auth_atom_id", "_atom_site.pdbx_PDB_model_num", }; size_t AtomSiteIndex(absl::string_view atom_site) { return std::distance(std::begin(kAtomSiteSortOrder), absl::c_find(kAtomSiteSortOrder, atom_site)); } struct AtomSiteOrder { bool operator()(absl::string_view lhs, absl::string_view rhs) const { auto lhs_index = AtomSiteIndex(lhs); auto rhs_index = AtomSiteIndex(rhs); return std::tie(lhs_index, lhs) < std::tie(rhs_index, rhs); } }; class Column { public: Column(absl::string_view key, const std::vector* values) : key_(key), values_(values) { int max_value_length = 0; for (size_t i = 0; i < values->size(); ++i) { absl::string_view value = (*values)[i]; if (IsTrivialToken(value)) { // Shortcut for the most common cases where no quoting/multiline needed. max_value_length = std::max(max_value_length, value.size()); continue; } else if (IsMultiLineToken(value)) { values_with_newlines_.insert(i); } else { absl::string_view quote = GetEscapeQuote(value); if (!quote.empty()) { values_with_quotes_[i] = quote; } max_value_length = std::max(max_value_length, value.size() + quote.size() * 2); } } max_value_length_ = max_value_length; } absl::string_view key() const { return key_; } const std::vector* values() const { return values_; } int max_value_length() const { return max_value_length_; } bool has_newlines(size_t index) const { return values_with_newlines_.contains(index); } absl::string_view quote(size_t index) const { if (auto it = values_with_quotes_.find(index); it != values_with_quotes_.end()) { return it->second; } return ""; } private: absl::string_view key_; const std::vector* values_; int max_value_length_; // Values with newlines or quotes are very rare in a typical CIF file. absl::flat_hash_set values_with_newlines_; absl::flat_hash_map values_with_quotes_; }; struct GroupedKeys { std::vector grouped_columns; int max_key_length; int value_size; }; absl::Status CheckLoopColumnSizes(int num_loop_keys, int num_loop_values) { if ((num_loop_keys > 0) && (num_loop_values % num_loop_keys != 0)) { return absl::InvalidArgumentError(absl::StrFormat( "The number of values (%d) in a loop is not a multiple of the " "number of the loop's columns (%d)", num_loop_values, num_loop_keys)); } return absl::OkStatus(); } } // namespace absl::StatusOr CifDict::FromString(absl::string_view cif_string) { CifDict::Dict cif; bool loop_flag = false; absl::string_view key; HeapStrings heap_strings; auto tokens = TokenizeInternal(cif_string, &heap_strings); if (!tokens.ok()) { return tokens.status(); } if (tokens->empty()) { return absl::InvalidArgumentError("The CIF file must not be empty."); } // The first token should be data_XXX. Split into key = data, value = XXX. absl::string_view first_token = tokens->front(); if (!absl::ConsumePrefix(&first_token, "data_")) { return absl::InvalidArgumentError( "The CIF file does not start with the data_ field."); } if (first_token.empty()) { return absl::InvalidArgumentError( "The CIF file does not contain a data block name."); } cif["data_"].emplace_back(first_token); // Counters for CIF loop_ regions. int loop_token_index = 0; int num_loop_keys = 0; // Loops have usually O(10) columns but could have up to O(10^6) rows. It is // therefore wasteful to look up the cif vector where to add a loop value // since that means doing `columns * rows` map lookups. If we save pointers to // these loop column fields instead, we need only 1 cif lookup per column. std::vector*> loop_column_values; // Skip the first element since we already processed it above. for (auto token_itr = tokens->begin() + 1; token_itr != tokens->end(); ++token_itr) { auto token = *token_itr; if (absl::EqualsIgnoreCase(token, "loop_")) { // A new loop started, check the previous loop and get rid of its data. absl::Status loop_status = CheckLoopColumnSizes(num_loop_keys, loop_token_index); if (!loop_status.ok()) { return loop_status; } loop_flag = true; loop_column_values.clear(); loop_token_index = 0; num_loop_keys = 0; continue; } else if (loop_flag) { // The second condition checks we are in the first column. Some mmCIF // files (e.g. 4q9r) have values in later columns starting with an // underscore and we don't want to read these as keys. int token_column_index = num_loop_keys == 0 ? 0 : loop_token_index % num_loop_keys; if (token_column_index == 0 && !token.empty() && token[0] == '_') { if (loop_token_index > 0) { // We are out of the loop. loop_flag = false; } else { // We are in the keys (column names) section of the loop. auto [it, inserted] = cif.try_emplace(token); if (!inserted) { return absl::InvalidArgumentError( absl::StrCat("Duplicate loop key: '", token, "'")); } auto& columns = it->second; columns.clear(); // Heuristic: _atom_site is typically the largest table in an mmCIF // with ~16 columns. Make sure we reserve enough space for its values. if (absl::StartsWith(token, "_atom_site.")) { columns.reserve(tokens->size() / 20); } // Save the pointer to the loop column values. loop_column_values.push_back(&columns); num_loop_keys += 1; continue; } } else { // We are in the values section of the loop. We have a pointer to the // loops' values, add the new token in there. if (token_column_index >= loop_column_values.size()) { return absl::InvalidArgumentError( absl::StrCat("Too many columns at: '", token, "' at column index: ", token_column_index, " expected at most: ", loop_column_values.size())); } loop_column_values[token_column_index]->emplace_back(token); loop_token_index++; continue; } } if (key.empty()) { key = token; if (!absl::StartsWith(key, "_")) { return absl::InvalidArgumentError( absl::StrCat("Key '", key, "' does not start with an underscore.")); } } else { auto [it, inserted] = cif.try_emplace(key); if (!inserted) { return absl::InvalidArgumentError( absl::StrCat("Duplicate key: '", key, "'")); } (it->second).emplace_back(token); key = ""; } } absl::Status loop_status = CheckLoopColumnSizes(num_loop_keys, loop_token_index); if (!loop_status.ok()) { return loop_status; } return CifDict(std::move(cif)); } absl::StatusOr CifDict::ToString() const { std::string output; absl::string_view data_name; // Check that the data_ field exists. if (auto name_it = (*dict_).find("data_"); name_it == (*dict_).end() || name_it->second.empty()) { return absl::InvalidArgumentError( "The CIF must contain a valid name for this data block in the special " "data_ field."); } else { data_name = name_it->second.front(); } if (absl::c_any_of(data_name, [](char i) { return absl::ascii_isspace(i); })) { return absl::InvalidArgumentError(absl::StrFormat( "The CIF data block name must not contain any whitespace characters, " "got '%s'.", data_name)); } absl::StrAppend(&output, "data_", data_name, "\n#\n"); // Group keys by their prefix. Use btree_map to iterate in alphabetical order, // but with some keys being placed at the end (e.g. _atom_site). absl::btree_map grouped_keys; for (const auto& [key, values] : *dict_) { if (key == "data_") { continue; // Skip the special data_ key, we are already done with it. } const std::pair key_parts = absl::StrSplit(key, absl::MaxSplits('.', 1)); const absl::string_view key_prefix = key_parts.first; auto [it, inserted] = grouped_keys.emplace(key_prefix, GroupedKeys{}); GroupedKeys& grouped_key = it->second; grouped_key.grouped_columns.push_back(Column(key, &values)); if (inserted) { grouped_key.max_key_length = key.length(); grouped_key.value_size = values.size(); } else { grouped_key.max_key_length = std::max(key.length(), grouped_key.max_key_length); if (grouped_key.value_size != values.size()) { return absl::InvalidArgumentError( absl::StrFormat("Values for key %s have different length (%d) than " "the other values with the same key prefix (%d).", key, values.size(), grouped_key.value_size)); } } } for (auto& [key_prefix, group_info] : grouped_keys) { if (key_prefix == "_atom_site") { // Make sure we sort the _atom_site loop in the standard way. absl::c_sort(group_info.grouped_columns, [](const Column& lhs, const Column& rhs) { return AtomSiteOrder{}(lhs.key(), rhs.key()); }); } else { // Make the key ordering within a key group deterministic. absl::c_sort(group_info.grouped_columns, [](const Column& lhs, const Column& rhs) { return lhs.key() < rhs.key(); }); } // Force `_atom_site` field to always be a loop. This resolves issues with // third party mmCIF parsers such as OpenBabel which always expect a loop // even when there is only a single atom present. if (group_info.value_size == 1 && key_prefix != "_atom_site") { // Plain key-value pairs, output them as they are. for (const Column& grouped_column : group_info.grouped_columns) { int width = group_info.max_key_length + 1; size_t start_pos = output.size(); output.append(width, ' '); auto out_it = output.begin() + start_pos; absl::c_copy(grouped_column.key(), out_it); // Append the value, handle multi-line/quoting. absl::string_view value = grouped_column.values()->front(); if (grouped_column.has_newlines(0)) { absl::StrAppend(&output, "\n;", value, "\n;\n"); // Multi-line value. } else { const absl::string_view quote_char = grouped_column.quote(0); absl::StrAppend(&output, quote_char, value, quote_char, "\n"); } } } else { // CIF loop. Output the column names, then the rows with data. absl::StrAppend(&output, "loop_\n"); for (Column& grouped_column : group_info.grouped_columns) { absl::StrAppend(&output, grouped_column.key(), "\n"); } // Write the loop values, line by line. This is the most expensive part // since this path is taken to write the entire atom site table which has // about 20 columns, but thousands of rows. for (int i = 0; i < group_info.value_size; i++) { for (int column_index = 0; column_index < group_info.grouped_columns.size(); ++column_index) { const Column& grouped_column = group_info.grouped_columns[column_index]; const absl::string_view value = (*grouped_column.values())[i]; if (grouped_column.has_newlines(i)) { // Multi-line. This is very rarely taken path. if (column_index == 0) { // No extra newline before leading ;, already inserted. absl::StrAppend(&output, ";", value, "\n;\n"); } else if (column_index == group_info.grouped_columns.size() - 1) { // No extra newline after trailing ;, will be inserted. absl::StrAppend(&output, "\n;", value, "\n;"); } else { absl::StrAppend(&output, "\n;", value, "\n;\n"); } } else { size_t start_pos = output.size(); output.append(grouped_column.max_value_length() + 1, ' '); auto out_it = output.begin() + start_pos; absl::string_view quote = grouped_column.quote(i); if (!quote.empty()) { out_it = absl::c_copy(quote, out_it); out_it = absl::c_copy(value, out_it); absl::c_copy(quote, out_it); } else { absl::c_copy(value, out_it); } } } absl::StrAppend(&output, "\n"); } } absl::StrAppend(&output, "#\n"); // Comment token after every key group. } return output; } absl::StatusOr< std::vector>> CifDict::ExtractLoopAsList(absl::string_view prefix) const { std::vector column_names; std::vector> column_data; for (const auto& element : *dict_) { if (absl::StartsWith(element.first, prefix)) { column_names.emplace_back(element.first); auto& cells = column_data.emplace_back(); cells.insert(cells.begin(), element.second.begin(), element.second.end()); } } // Make sure all columns have the same number of rows. const std::size_t num_rows = column_data.empty() ? 0 : column_data[0].size(); for (const auto& column : column_data) { if (column.size() != num_rows) { return absl::InvalidArgumentError(absl::StrCat( GetDataName(), ": Columns do not have the same number of rows for prefix: '", prefix, "'. One possible reason could be not including the trailing dot, " "e.g. '_atom_site.'.")); } } std::vector> result; result.reserve(num_rows); CHECK_EQ(column_names.size(), column_data.size()); for (std::size_t row_index = 0; row_index < num_rows; ++row_index) { auto& row_dict = result.emplace_back(); row_dict.reserve(column_names.size()); for (int col_index = 0; col_index < column_names.size(); ++col_index) { row_dict[column_names[col_index]] = column_data[col_index][row_index]; } } return result; } absl::StatusOr>> CifDict::ExtractLoopAsDict(absl::string_view prefix, absl::string_view index) const { if (!absl::StartsWith(index, prefix)) { return absl::InvalidArgumentError( absl::StrCat(GetDataName(), ": The loop index '", index, "' must start with the loop prefix '", prefix, "'.")); } absl::flat_hash_map> result; auto loop_as_list = ExtractLoopAsList(prefix); if (!loop_as_list.ok()) { return loop_as_list.status(); } result.reserve(loop_as_list->size()); for (auto& entry : *loop_as_list) { if (const auto it = entry.find(index); it != entry.end()) { result[it->second] = entry; } else { return absl::InvalidArgumentError(absl::StrCat( GetDataName(), ": The index column '", index, "' could not be found in the loop with prefix '", prefix, "'.")); } } return result; } absl::StatusOr> Tokenize( absl::string_view cif_string) { HeapStrings heap_strings; auto tokens = TokenizeInternal(cif_string, &heap_strings); if (!tokens.ok()) { return tokens.status(); } return std::vector(tokens->begin(), tokens->end()); } absl::StatusOr> SplitLine( absl::string_view line) { std::vector tokens; if (!SplitLineInline(line, &tokens)) { return absl::InvalidArgumentError( absl::StrCat("Line ended with quote open: ", line)); } return tokens; } absl::StatusOr> ParseMultiDataCifDict( absl::string_view cif_string) { absl::flat_hash_map mapping; constexpr absl::string_view delimitor = "data_"; // Check cif_string starts with correct offset. if (!cif_string.empty() && !absl::StartsWith(cif_string, delimitor)) { return absl::InvalidArgumentError( "Invalid format. MultiDataCifDict must start with 'data_'"); } for (absl::string_view data_block : absl::StrSplit(cif_string, delimitor, absl::SkipEmpty())) { absl::string_view block_with_delimitor( data_block.data() - delimitor.size(), data_block.size() + delimitor.size()); absl::StatusOr parsed_block = CifDict::FromString(block_with_delimitor); if (!parsed_block.ok()) { return parsed_block.status(); } absl::string_view data_name = parsed_block->GetDataName(); mapping[data_name] = *std::move(parsed_block); } return mapping; } } // namespace alphafold3