Merge junction beads when they are very close; removed cruft

segmentmodel.py

@@ -13,14 +13,12 @@ import types
 import pdb
 """
 TODO:
- - document
- - handle crossovers
-    - connections in the middle of a segment?
-    - merge beads at ends of connected helices?
+ - fix handling of crossovers for atomic representation
  - map to atomic representation
  - remove performance bottlenecks
  - test for large systems
  - assign sequence
+ - document
 """
 
 class Location():
@@ -190,8 +188,13 @@ class Segment(ConnectableElement, Group):
     def _generate_one_bead(self, contour_position, nts):
         raise NotImplementedError
 
-    def _assign_particles_to_locations(self):
-        raise NotImplementedError
+    def get_nearest_bead(self, contour_position):
+        if len(self.beads) < 1: return None
+        cs = np.array([b.contour_position for b in self.beads]) # TODO: cache
+        # TODO: include connected beads?
+        i = np.argmin((cs - contour_position)**2)
+
+        return self.beads[i]
 
     def get_all_consecutive_beads(self, number):
         assert(number >= 1)
@@ -349,7 +352,7 @@ class DoubleStrandedSegment(Segment):
 
 
     ## Convenience methods
-    ## TODO: add errors if incorrect connections are made
+    ## TODO: add errors if unrealistic connections are made
     def connect_start5(self, end3, type_="intrahelical", force_connection=False):
         if isinstance(end3, SingleStrandedSegment):
             end3 = end3.end3
@@ -439,17 +442,6 @@ class DoubleStrandedSegment(Segment):
         return bead
 
 
-
-    def _assign_particles_to_locations(self):
-        if self.start5.particle is None:
-            assert(self.beads[0].parent is not None)
-            self.start5.particle = self.beads[0]
-        if self.end3.particle is None:
-            assert(self.beads[-1].parent is not None)
-            self.end3.particle = self.beads[-1]
-        self.start3.particle =  self.start5.particle 
-        self.end5.particle   =  self.end3.particle
-
     def _generate_atomic(self, atomic_model):
         ...
     
@@ -520,14 +512,6 @@ class SingleStrandedSegment(Segment):
         self._add_bead(b)
         return b
 
-    def _assign_particles_to_locations(self):
-        if self.start.particle is None:
-            assert(self.beads[0].parent is not None)
-            self.start.particle = self.beads[0]
-        if self.end.particle is None:
-            assert(self.beads[-1].parent is not None)
-            self.end.particle = self.beads[-1]
-
     def _generate_atomic(self, atomic_model):
         ...
     
@@ -745,7 +729,7 @@ class SegmentModel(ArbdModel):
         for s in segments:
             s.local_twist = local_twist
             
-        """ Generate beads at junctions """
+        """ Generate beads at intrahelical junctions """
         if self.DEBUG: print( "Adding intrahelical beads at junctions" )
         ## Loop through all connections, generate beads at appropriate locations
         for c,A,B in self.get_connections("intrahelical"):
@@ -770,6 +754,7 @@ class SegmentModel(ArbdModel):
                 A.particle = B.particle = b
                 ...
 
+        """ Generate beads at other junctions """
         for c,A,B in self.get_connections(exclude="intrahelical"):
             s1,s2 = [l.container for l in (A,B)]
             if A.particle is not None and B.particle is not None:
@@ -777,11 +762,24 @@ class SegmentModel(ArbdModel):
             assert( A.particle is None )
             assert( B.particle is None )
 
-            ## TODO: offload the work here to s1
+            ## TODO: offload the work here to s1/s2 (?)
             a1,a2 = [l.address   for l in (A,B)]
-            ## TODO: if existing particle of same type is very nearby, use that instead
-            A.particle = s1._generate_one_bead(a1,0)
-            B.particle = s2._generate_one_bead(a2,0)
+
+            b = s1.get_nearest_bead(a1)
+            if b is not None and np.abs(b.contour_position-a1)*s1.num_nts < 1:
+                ## combine beads
+                b.contour_position = 0.5*(b.contour_position + a1) # avg position
+                A.particle = b
+            else:
+                A.particle = s1._generate_one_bead(a1,0)
+
+            b = s2.get_nearest_bead(a2)
+            if b is not None and np.abs(b.contour_position-a2)*s2.num_nts < 1:
+                ## combine beads
+                b.contour_position = 0.5*(b.contour_position + a2) # avg position
+                B.particle = b
+            else:
+                B.particle = s2._generate_one_bead(a2,0)
 
         """ Some tests """
         for c,A,B in self.get_connections("intrahelical"):
@@ -1021,13 +1019,4 @@ class SegmentModel(ArbdModel):
                         t0 = 180
                         pot = self.get_dihedral_potential(k,t0)
                         self.add_dihedral( u1,b1,b2,d2, pot )
-
-
-    # def get_bead(self, location):
-    #     if type(location.container) is not list:
-    #         s = self.segments.index(location.container)
-    #         s.get_bead(location.address)
-    #     else:
-    #         r
-    #         ...