Aleaverfay/kinforest from stencil

pytest.ini

@@ -2,7 +2,7 @@
 ext = .cpp,.hh,.h,.cu,.py,.yaml
 
 [pytest]
-addopts = --benchmark-disable --benchmark-columns=ops,mean,iqr
+# addopts = --benchmark-disable --benchmark-columns=ops,mean,iqr
 filterwarnings =
   ignore:(?s).*is not compatible with the compiler Pytorch.*:
   ignore:(?s).*Benchmark fixture was not used.*:

tmol/chemical/patched_chemdb.py

@@ -510,6 +510,7 @@ def do_patch(res, variant, resgraph, patchgraph, marked):
             icoors=res.icoors,
             properties=res.properties,
             chi_samples=res.chi_samples,
+            default_jump_connection_atom=res.default_jump_connection_atom,
         )
 
         # 1. remove atoms

tmol/chemical/restypes.py

@@ -155,6 +155,10 @@ def _setup_bond_indices(self):
         bond_array.flags.writeable = False
         return bond_array
 
+    @property
+    def n_conn(self):
+        return len(self.connections)
+
     # The index of the atom for a given inter-residue connection point
     connection_to_idx: Mapping[str, AtomIndex] = attr.ib()
 
@@ -259,6 +263,10 @@ def _setup_ordered_torsions(self):
                 )
         return ordered_torsions
 
+    @property
+    def n_torsions(self):
+        return self.ordered_torsions.shape[0]
+
     path_distance: numpy.ndarray = attr.ib()
 
     @path_distance.default
@@ -448,6 +456,12 @@ def _setup_icoors_geom(self):
     def compute_ideal_coords(self):
         return build_coords_from_icoors(self.icoors_ancestors, self.icoors_geom)
 
+    default_jump_connection_atom_index: int = attr.ib()
+
+    @default_jump_connection_atom_index.default
+    def get_default_jump_connection_atom_index(self):
+        return self.atom_to_idx[self.default_jump_connection_atom]
+
 
 @attr.s(auto_attribs=True)
 class ResidueTypeSet:

tmol/database/chemical/__init__.py

@@ -136,6 +136,10 @@ class RawResidueType:
     icoors: Tuple[Icoor, ...]
     properties: ChemicalProperties
     chi_samples: Tuple[ChiSamples, ...]
+    default_jump_connection_atom: str
+
+    def atom_name(self, index):
+        return self.atoms[index].name
 
 
 @attr.s(auto_attribs=True, frozen=True, slots=True)

tmol/database/default/chemical/chemical.yaml

@@ -143,6 +143,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  ARG
     base_name: ARG
     name3: ARG
@@ -277,6 +278,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  ASN
     base_name: ASN
     name3: ASN
@@ -373,6 +375,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  ASP
     base_name: ASP
     name3: ASP
@@ -461,6 +464,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  CYS
     base_name: CYS
     name3: CYS
@@ -549,6 +553,7 @@ residues:
     - chi_dihedral: chi2
       samples: [60, 180, 300]
       expansions: []
+    default_jump_connection_atom: CA
   - name:  CYD
     base_name: CYD
     name3: CYS
@@ -631,6 +636,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  GLN
     base_name: GLN
     name3: GLN
@@ -739,6 +745,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  GLU
     base_name: GLU
     name3: GLU
@@ -839,6 +846,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  GLY
     base_name: GLY
     name3: GLY
@@ -910,6 +918,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  HIS
     base_name: HIS
     name3: HIS
@@ -1016,6 +1025,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  HIS_D
     base_name: HIS_D
     name3: HIS
@@ -1121,6 +1131,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  ILE
     base_name: ILE
     name3: ILE
@@ -1236,6 +1247,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  LEU
     base_name: LEU
     name3: LEU
@@ -1351,6 +1363,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  LYS
     base_name: LYS
     name3: LYS
@@ -1480,6 +1493,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  MET
     base_name: MET
     name3: MET
@@ -1589,6 +1603,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  PHE
     base_name: PHE
     name3: PHE
@@ -1702,6 +1717,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  PRO
     base_name: PRO
     name3: PRO
@@ -1802,6 +1818,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  SER
     base_name: SER
     name3: SER
@@ -1890,6 +1907,7 @@ residues:
     - chi_dihedral: chi2
       samples: [0, 20, 40, 60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280, 300, 320, 340]
       expansions: []
+    default_jump_connection_atom: CA
   - name:  THR
     base_name: THR
     name3: THR
@@ -1988,6 +2006,7 @@ residues:
     - chi_dihedral: chi2
       samples: [0, 20, 40, 60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280, 300, 320, 340]
       expansions: []
+    default_jump_connection_atom: CA
   - name:  TRP
     base_name: TRP
     name3: TRP
@@ -2114,6 +2133,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  TYR
     base_name: TYR
     name3: TYR
@@ -2235,6 +2255,7 @@ residues:
     - chi_dihedral: chi3
       samples: [0, 180]
       expansions: [20]
+    default_jump_connection_atom: CA
   - name:  VAL
     base_name: VAL
     name3: VAL
@@ -2339,6 +2360,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: CA
   - name:  HOH
     base_name: HOH
     name3: HOH
@@ -2374,6 +2396,7 @@ residues:
         pH: 7
       virtual: []
     chi_samples: []
+    default_jump_connection_atom: O
 
 variants:
   - name:  CarboxyTerminus

tmol/extern/moderngpu/cta_scan.hxx

@@ -2,14 +2,10 @@
 #pragma once
 #include "loadstore.hxx"
 #include "intrinsics.hxx"
+#include "scan_types.hxx"
 
 BEGIN_MGPU_NAMESPACE
 
-enum scan_type_t {
-  scan_type_exc,
-  scan_type_inc
-};
-
 template<typename type_t, int vt = 0, bool is_array = (vt > 0)>
 struct scan_result_t {
   type_t scan;
@@ -32,7 +28,7 @@ struct cta_scan_t {
     struct { type_t threads[nt], warps[num_warps]; };
   };
 
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 300  
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 300
 
   //////////////////////////////////////////////////////////////////////////////
   // Optimized CTA scan code that uses warp shfl intrinsics.
@@ -41,7 +37,7 @@ struct cta_scan_t {
 
   template<typename op_t = plus_t<type_t> >
   MGPU_DEVICE scan_result_t<type_t>
-  scan(int tid, type_t x, storage_t& storage, int count = nt, op_t op = op_t(), 
+  scan(int tid, type_t x, storage_t& storage, int count = nt, op_t op = op_t(),
     type_t init = type_t(), scan_type_t type = scan_type_exc) const {
 
     int warp = tid / warp_size;
@@ -61,7 +57,7 @@ struct cta_scan_t {
     __syncthreads();
 
     // Scan the warp reductions.
-    if(tid < num_warps) { 
+    if(tid < num_warps) {
       type_t cta_scan = storage.warps[tid];
       iterate<s_log2(num_warps)>([&](int pass) {
         cta_scan = shfl_up_op(cta_scan, 1<< pass, op, num_warps);
@@ -78,10 +74,10 @@ struct cta_scan_t {
     if(warp > 0) scan = op(scan, storage.warps[warp - 1]);
 
     type_t reduction = storage.warps[div_up(count, warp_size) - 1];
-    
-    scan_result_t<type_t> result { 
-      tid < count ? scan : reduction, 
-      reduction 
+
+    scan_result_t<type_t> result {
+      tid < count ? scan : reduction,
+      reduction
     };
     __syncthreads();
 
@@ -91,11 +87,11 @@ struct cta_scan_t {
 #else
 
   //////////////////////////////////////////////////////////////////////////////
-  // Standard CTA scan code that does not use shfl intrinsics. 
+  // Standard CTA scan code that does not use shfl intrinsics.
 
   template<typename op_t = plus_t<type_t> >
-  MGPU_DEVICE scan_result_t<type_t> 
-  scan(int tid, type_t x, storage_t& storage, int count = nt, op_t op = op_t(), 
+  MGPU_DEVICE scan_result_t<type_t>
+  scan(int tid, type_t x, storage_t& storage, int count = nt, op_t op = op_t(),
     type_t init = type_t(), scan_type_t type = scan_type_exc) const {
 
     int first = 0;
@@ -113,7 +109,7 @@ struct cta_scan_t {
 
     scan_result_t<type_t> result;
     result.reduction = storage.data[first + count - 1];
-    result.scan = (tid < count) ? 
+    result.scan = (tid < count) ?
       (scan_type_inc == type ? x :
         (tid ? storage.data[first + tid - 1] : init)) :
       result.reduction;
@@ -122,16 +118,16 @@ struct cta_scan_t {
     return result;
   }
 
-#endif  
+#endif
 
   //////////////////////////////////////////////////////////////////////////////
-  // CTA vectorized scan. Accepts multiple values per thread and adds in 
+  // CTA vectorized scan. Accepts multiple values per thread and adds in
   // optional global carry-in.
 
   template<int vt, typename op_t = plus_t<type_t> >
   MGPU_DEVICE scan_result_t<type_t, vt>
-  scan(int tid, array_t<type_t, vt> x, storage_t& storage, 
-    type_t carry_in = type_t(), bool use_carry_in = false, 
+  scan(int tid, array_t<type_t, vt> x, storage_t& storage,
+    type_t carry_in = type_t(), bool use_carry_in = false,
     int count = nt, op_t op = op_t(), type_t init = type_t(),
     scan_type_t type = scan_type_exc) const {
 
@@ -143,14 +139,14 @@ struct cta_scan_t {
     } else {
       iterate<vt>([&](int i) {
         int index = vt * tid + i;
-        x[i] = i ? 
+        x[i] = i ?
           ((index < count) ? op(x[i], x[i - 1]) : x[i - 1]) :
           (x[i] = (index < count) ? x[i] : init);
       });
     }
 
     // Scan the thread-local reductions for a carry-in for each thread.
-    scan_result_t<type_t> result = scan(tid, x[vt - 1], storage, 
+    scan_result_t<type_t> result = scan(tid, x[vt - 1], storage,
       div_up(count, vt), op, init, scan_type_exc);
 
     // Perform the scan downsweep and add both the global carry-in and the
@@ -185,7 +181,7 @@ struct cta_scan_t<nt, bool> {
     int warps[num_warps];
   };
 
-  MGPU_DEVICE scan_result_t<int> scan(int tid, bool x, 
+  MGPU_DEVICE scan_result_t<int> scan(int tid, bool x,
     storage_t& storage) const {
 
     // Store the bit totals for each warp.
@@ -207,7 +203,7 @@ struct cta_scan_t<nt, bool> {
     }
     __syncthreads();
 #else
-    
+
     if(0 == tid) {
       // Inclusive scan of partial reductions..
       int scan = 0;
@@ -217,7 +213,7 @@ struct cta_scan_t<nt, bool> {
     }
     __syncthreads();
 
-#endif    
+#endif
 
     int scan = ((warp > 0) ? storage.warps[warp - 1] : 0) +
       popc(bfe(bits, 0, lane));