static int ks_assign_hydrogens(const float* xyz, const int* nco_indices, const int n_residues, float *hcoords, int* skip) /* Assign hydrogen atom coordinates */ { int ri, pc_index, po_index; __m128 pc, po, r_co, r_h, r_n, norm_r_co; __m128 tenth = _mm_set1_ps(0.1f); r_n = load_float3(xyz + 3*nco_indices[0]); store_float3(hcoords, r_n); hcoords += 3; for (ri = 1; ri < n_residues; ri++) { if (!skip[ri]) { pc_index = nco_indices[3*(ri-1) + 1]; po_index = nco_indices[3*(ri-1) + 2]; pc = load_float3(xyz + 3*pc_index); po = load_float3(xyz + 3*po_index); r_co = _mm_sub_ps(pc, po); r_n = load_float3(xyz + 3*nco_indices[3*ri + 0]); norm_r_co = _mm_mul_ps(r_co, _mm_rsqrt_ps(_mm_dp_ps2(r_co, r_co, 0xFF))); r_h = _mm_add_ps(r_n, _mm_mul_ps(tenth, norm_r_co)); store_float3(hcoords, r_h); } hcoords += 3; } return 1; }
static void asa_frame(const float* frame, const int n_atoms, const float* atom_radii, const float* sphere_points, const int n_sphere_points, int* neighbor_indices, float* centered_sphere_points, float* areas) { /*// Calculate the accessible surface area of each atom in a single snapshot // // Parameters // ---------- // frame : 2d array, shape=[n_atoms, 3] // The coordinates of the nuclei // n_atoms : int // the major axis length of frame // atom_radii : 1d array, shape=[n_atoms] // the van der waals radii of the atoms PLUS the probe radius // sphere_points : 2d array, shape=[n_sphere_points, 3] // a bunch of uniformly distributed points on a sphere // n_sphere_points : int // the number of sphere points // centered_sphere_points : WORK BUFFER 2d array, shape=[n_sphere_points, 3] // empty memory that intermediate calculations can be stored in // neighbor_indices : WORK BUFFER 2d array, shape=[n_atoms] // empty memory that intermediate calculations can be stored in // NOTE: the point of these work buffers is that if we want to call // this function repreatedly, its more efficient not to keep re-mallocing // these work buffers, but instead just reuse them. // areas : 1d array, shape=[n_atoms] // the output buffer to place the results in -- the surface area of each // atom */ int i, j, k, k_prime; __m128 r, r_i, r_j, r_ij, atom_radius_i, atom_radius_j, radius_cutoff; __m128 radius_cutoff2, sp, r_jk, r2; int n_neighbor_indices, is_accessible, k_closest_neighbor; float constant = 4.0 * M_PI / n_sphere_points; for (i = 0; i < n_atoms; i++) { atom_radius_i = _mm_set1_ps(atom_radii[i]); r_i = load_float3(frame+i*3); /* Get all the atoms close to atom `i` */ n_neighbor_indices = 0; for (j = 0; j < n_atoms; j++) { if (i == j) { continue; } r_j = load_float3(frame+j*3); r_ij = _mm_sub_ps(r_i, r_j); atom_radius_j = _mm_set1_ps(atom_radii[j]); /* Look for atoms `j` that are nearby atom `i` */ radius_cutoff = _mm_add_ps(atom_radius_i, atom_radius_j); radius_cutoff2 = _mm_mul_ps(radius_cutoff, radius_cutoff); r2 = _mm_dp_ps(r_ij, r_ij, 0x7F); if (_mm_extract_epi16(CAST__M128I(_mm_cmplt_ps(r2, radius_cutoff2)), 0)) { neighbor_indices[n_neighbor_indices] = j; n_neighbor_indices++; } if (_mm_extract_epi16(CAST__M128I(_mm_cmplt_ps(r2, _mm_set1_ps(1e-10))), 0)) { printf("ERROR: THIS CODE IS KNOWN TO FAIL WHEN ATOMS ARE VIRTUALLY"); printf("ON TOP OF ONE ANOTHER. YOU SUPPLIED TWO ATOMS %f", _mm_cvtss_f32(r)); printf("APART. QUITTING NOW"); exit(1); } } /* Center the sphere points on atom i */ for (j = 0; j < n_sphere_points; j++) { sp = _mm_add_ps(r_i, _mm_mul_ps(atom_radius_i, load_float3(sphere_points + 3*j))); store_float3(centered_sphere_points + 3*j, sp); } /* Check if each of these points is accessible */ k_closest_neighbor = 0; for (j = 0; j < n_sphere_points; j++) { is_accessible = 1; r_j = load_float3(centered_sphere_points + 3*j); /* iterate through the sphere points by cycling through them */ /* in a circle, starting with k_closest_neighbor and then wrapping */ /* around */ for (k = k_closest_neighbor; k < n_neighbor_indices + k_closest_neighbor; k++) { k_prime = k % n_neighbor_indices; r = _mm_set1_ps(atom_radii[neighbor_indices[k_prime]]); r_jk = _mm_sub_ps(r_j, load_float3(frame+3*neighbor_indices[k_prime])); if (_mm_extract_epi16(CAST__M128I(_mm_cmplt_ps(_mm_dp_ps(r_jk, r_jk, 0xFF), _mm_mul_ps(r, r))), 0)) { k_closest_neighbor = k; is_accessible = 0; break; } } if (is_accessible) { areas[i]++; } } areas[i] *= constant * (atom_radii[i])*(atom_radii[i]); } }