segmentmodel.py 49 KB
Newer Older
1
2
import numpy as np
from arbdmodel import PointParticle, ParticleType, Group, ArbdModel
3
from coords import rotationAboutAxis, quaternion_from_matrix, quaternion_to_matrix
4
5
6
7
from nonbonded import *
from copy import copy, deepcopy
from nbPot import nbDnaScheme

cmaffeo2's avatar
cmaffeo2 committed
8
9
from scipy.special import erf
import scipy.optimize as opt
10
from scipy import interpolate
cmaffeo2's avatar
cmaffeo2 committed
11

cmaffeo2's avatar
cmaffeo2 committed
12
13
14
from CanonicalNucleotideAtoms import canonicalNtFwd, canonicalNtRev, seqComplement
from CanonicalNucleotideAtoms import enmTemplateHC, enmTemplateSQ, enmCorrectionsHC

15
import pdb
16
"""
cmaffeo2's avatar
cmaffeo2 committed
17
TODO:
cmaffeo2's avatar
cmaffeo2 committed
18
 + fix handling of crossovers for atomic representation
cmaffeo2's avatar
cmaffeo2 committed
19
20
21
22
23
24
 + map to atomic representation
    - add nicks
    - shrink ssDNA
    - shrink dsDNA backbone
    - make orientation continuous
    - sequence
25
26
27
 - remove performance bottlenecks
 - test for large systems
 - assign sequence
28
 - document
29
30
31
32
33
"""

class Location():
    """ Site for connection within an object """
    def __init__(self, container, address, type_):
34
        ## TODO: remove cyclic references(?)
35
        self.container = container
36
        self.address = address  # represents position along contour length in segments
cmaffeo2's avatar
cmaffeo2 committed
37
        # assert( type_ in ("end3","end5") ) # TODO remove or make conditional
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
        self.type_ = type_
        self.particle = None

class Connection():
    """ Abstract base class for connection between two elements """
    def __init__(self, A, B, type_ = None):
        assert( isinstance(A,Location) )
        assert( isinstance(B,Location) )
        self.A = A
        self.B = B
        self.type_ = type_
        
# class ConnectableElement(Transformable):
class ConnectableElement():
    """ Abstract base class """
    def __init__(self, connections=[]):
        self.connections = connections

cmaffeo2's avatar
cmaffeo2 committed
56
    def get_connections_and_locations(self, type_=None, exclude=[]):
57
58
59
60
        """ Returns a list with each entry of the form:
            connection, location_in_self, location_in_other """
        ret = []
        for c in self.connections:
cmaffeo2's avatar
cmaffeo2 committed
61
            if type_ is None or c.type_ == type_ and type_ not in exclude:
62
                if   c.A.container is self:
63
                    ret.append( [c, c.A, c.B] )
64
                elif c.B.container is self:
65
66
67
68
69
                    ret.append( [c, c.B, c.A] )
                else:
                    raise Exception("Object contains connection that fails to refer to object")
        return ret

70
71
72
    def _connect(self, other, connection):
        self.connections.append(connection)
        other.connections.append(connection)
73
74
    # def _find_connections(self, loc):
    #     return [c for c in self.connections if c.A == loc or c.B == loc]
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

class SegmentParticle(PointParticle):
    def __init__(self, type_, position, name="A", segname="A", **kwargs):
        self.contour_position = None
        PointParticle.__init__(self, type_, position, name=name, segname=segname, **kwargs)
        self.intrahelical_neighbors = []
        self.other_neighbors = []

    def get_contour_position(self,seg):
        assert( isinstance(seg,Segment) )
        if seg == self.parent:
            return self.contour_position
        else:
            for c,A,B in self.parent.get_connections_and_locations():
                if A.particle is self and B.container is seg:
                    nt = np.abs( (self.contour_position - A.address)*A.container.num_nts )
                    if B.address < 0.5:
                        return B.address-nt/seg.num_nts
                    else:
                        return B.address+nt/seg.num_nts
cmaffeo2's avatar
cmaffeo2 committed
95
            ## ERROR
96
97
98
99
100
101
102
103
104
105
            print("")
            for c,A,B in self.parent.get_connections_and_locations():
                print("  ",c.type_)
                print(A,B)
                print(A.particle,self)
                print(B.container,seg)
            print("")
            import pdb
            pdb.set_trace()
            raise Exception("Did not find location for particle {} in Segment {}".format(self,seg))
106
107
108
109
110
111
112
113
114
115
116
117
       
class Segment(ConnectableElement, Group):

    """ Base class that describes a segment of DNA. When built from
    cadnano models, should not span helices """

    """Define basic particle types"""
    dsDNA_particle = ParticleType("D",
                                  diffusivity = 43.5,
                                  mass = 300,
                                  radius = 3,                 
                              )
cmaffeo2's avatar
cmaffeo2 committed
118
119
120
121
122
    orientation_particle = ParticleType("O",
                                        diffusivity = 100,
                                        mass = 300,
                                        radius = 1,
                                    )
123

cmaffeo2's avatar
cmaffeo2 committed
124
    # orientation_bond = HarmonicBond(10,2)
125
    orientation_bond = HarmonicBond(30,1.5, rRange = (0,500) )
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

    ssDNA_particle = ParticleType("S",
                                  diffusivity = 43.5,
                                  mass = 150,
                                  radius = 3,                 
                              )

    def __init__(self, name, num_nts, 
                 start_position = np.array((0,0,0)),
                 end_position = None, 
                 segment_model = None):

        Group.__init__(self, name, children=[])
        ConnectableElement.__init__(self, connections=[])

cmaffeo2's avatar
cmaffeo2 committed
141
142
143
144
145
        self.start_orientation = None
        self.twist_per_nt = 0

        self.beads = [c for c in self.children] # self.beads will not contain orientation beads

146
147
148
        self._bead_model_generation = 0    # TODO: remove?
        self.segment_model = segment_model # TODO: remove?

cmaffeo2's avatar
cmaffeo2 committed
149
        self.num_nts = int(num_nts)
150
151
152
153
154
        if end_position is None:
            end_position = np.array((0,0,self.distance_per_nt*num_nts)) + start_position
        self.start_position = start_position
        self.end_position = end_position

155
156
157
158
159
        ## Set up interpolation for positions
        a = np.array([self.start_position,self.end_position]).T
        tck, u = interpolate.splprep( a, u=[0,1], s=0, k=1)
        self.position_spline_params = tck

160
161
162
163
164
    def clear_all(self):
        Group.clear_all(self)  # TODO: use super?
        self.beads = []
        for c,loc,other in self.get_connections_and_locations():
            loc.particle = None
165
166
167
168
169
170
171
172

    def contour_to_position(self,s):
        p = interpolate.splev( s, self.position_spline_params )
        if len(p) > 1: p = np.array(p).T
        return p

    def contour_to_tangent(self,s):
        t = interpolate.splev( s, self.position_spline_params, der=1 )
173
174
        t = (t / np.linalg.norm(t,axis=0))
        return t.T
175
176
177
178
179
180
181
        

    def contour_to_orientation(self,s):
        if self.start_orientation is not None:
            # axis = self.start_orientation.dot( np.array((0,0,1)) )
            if self.quaternion_spline_params is None:
                axis = self.contour_to_tangent(s)
182
                orientation = rotationAboutAxis( axis, s*self.twist_per_nt*self.num_nts, normalizeAxis=True )
183
184
185
186
187
                ## TODO: ensure this is correct
                # orientation = self.start_orientation.dot(orientation) # .dot( self.start_orientation )
                orientation = orientation.dot( self.start_orientation )
            else:
                q = interpolate.splev(s, self.quaternion_spline_params)
cmaffeo2's avatar
cmaffeo2 committed
188
                q = q/np.linalg.norm(q)
189
                orientation = quaternion_to_matrix(q)
cmaffeo2's avatar
cmaffeo2 committed
190
                
191
192
193
194
        else:
            orientation = None
        return orientation

cmaffeo2's avatar
cmaffeo2 committed
195
196
197
198
    def get_contour_sorted_connections_and_locations(self):
        sort_fn = lambda c: c[1].address
        cl = self.get_connections_and_locations()
        return sorted(cl, key=sort_fn)
199

cmaffeo2's avatar
cmaffeo2 committed
200
201
202
    def _get_num_beads(self, max_basepairs_per_bead, max_nucleotides_per_bead ):
        raise NotImplementedError

203
    def _generate_one_bead(self, contour_position, nts):
204
205
        raise NotImplementedError

cmaffeo2's avatar
cmaffeo2 committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
    def _generate_atomic_nucleotide(self, contour_position, is_fwd, seq):
        """ Seq should include modifications like 5T, T3 Tsinglet; direction matters too """

        # print("Generating nucleotide at {}".format(contour_position))
        
        pos = self.contour_to_position(contour_position)
        if self.local_twist:
            orientation = self.contour_to_orientation(contour_position)
            ## TODO: move this code (?)
            if orientation is None:
                axis = self.contour_to_tangent(contour_position)
                angleVec = np.array([1,0,0])
                if axis.dot(angleVec) > 0.9: angleVec = np.array([0,1,0])
                angleVec = angleVec - angleVec.dot(axis)*axis
                angleVec = angleVec/np.linalg.norm(angleVec)
                y = np.cross(axis,angleVec)
                orientation = np.array([angleVec,y,axis]).T
                ## TODO: improve placement of ssDNA
                # rot = rotationAboutAxis( axis, contour_position*self.twist_per_nt*self.num_nts, normalizeAxis=True )
                # orientation = rot.dot(orientation)
            else:
                orientation = orientation
                            
        else:
            raise NotImplementedError

        # key = self.sequence
        # if self.ntAt5prime is None and self.ntAt3prime is not None: key = "5"+key
        # if self.ntAt5prime is not None and self.ntAt3prime is None: key = key+"3"
        # if self.ntAt5prime is None and self.ntAt3prime is None: key = key+"singlet"

        key = seq

        ## TODO
        if not is_fwd:
            nt_dict = canonicalNtFwd
        else:
            nt_dict = canonicalNtRev
        atoms = nt_dict[ key ].duplicate() # TODO: clone
                        
        # print( atoms.orientation, orientation )
        # atoms.orientation = atoms.orientation.dot( orientation )
        atoms.orientation = orientation.dot(atoms.orientation)
        atoms.position = pos

        ## TODO: scale positions

        return atoms
        
    def get_5prime_locations(self):
        """ Returns tuple of contour_position and direction of strand
        True represents a strand whose 5-to-3 direction increases with contour
        """
        raise NotImplementedError

    def get_end_of_strand(self, contour_pos, is_fwd):

        ## connections to other segments
        cl = self.get_contour_sorted_connections_and_locations()
        if not is_fwd:          # reverse sorted list if not forward
            cl = cl[::-1]

        ## TODO include nicks
        ## Iterate through connection locations in segment
        for c,l,B in cl:
            if l.particle is None:
                pos = l.address
            else:
                pos = l.particle.get_contour_position()
                
            ## skip locations that are not in the correct direction
            if is_fwd:
                if pos <= contour_pos: continue
            elif pos >= contour_pos: continue

            ## Stop at every crossover
            if c.type_ == "crossover":
                if l.type_ == is_fwd:
                    print("   crossover -> cross")
                    return pos, B.container, B.address, B.type_
                else:
                    print("   crossover -> continue")
                    return pos, l.container, pos, is_fwd # break here, but continue strand

            if l.type_ ==  "end3":
                print("   end3") 
                assert( np.abs(B.address) < 1e-2 or np.abs(B.address-1) < 1e-2 )
                direction = B.address < 0.5
                return pos, B.container, B.address, direction
        if is_fwd:
            return 1, None, None, None
        else:
            return 0, None, None, None
        
300
301
302
303
304
305
306
    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]
307
308
309

    def get_all_consecutive_beads(self, number):
        assert(number >= 1)
cmaffeo2's avatar
cmaffeo2 committed
310
        ## Assume that consecutive beads in self.beads are bonded
311
        ret = []
cmaffeo2's avatar
cmaffeo2 committed
312
313
        for i in range(len(self.beads)-number+1):
            tmp = [self.beads[i+j] for j in range(0,number)]
314
            ret.append( tmp )
315
        return ret   
316

317
318
319
    def _add_bead(self,b,set_contour=False):
        if set_contour:
            b.contour_position = b.get_contour_position(self)
320
        
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
        # assert(b.parent is None)
        if b.parent is not None:
            b.parent.children.remove(b)
        b.parent = self
        self.children.append(b)
        self.beads.append(b) # don't add orientation bead
        if "orientation_bead" in b.__dict__: # TODO: think of a cleaner approach
            o = b.orientation_bead
            o.contour_position = b.contour_position
            if o.parent is not None:
                o.parent.children.remove(o)
            o.parent = self
            self.children.append(o)
            self.add_bond(b,o, Segment.orientation_bond, exclude=True)

    def _rebuild_children(self, new_children):
        # print("_rebuild_children on %s" % self.name)
        old_children = self.children
        old_beads = self.beads
        self.children = []
        self.beads = []
        for b in new_children:
            self.beads.append(b)
            self.children.append(b)
            if "orientation_bead" in b.__dict__: # TODO: think of a cleaner approach
                self.children.append(b.orientation_bead)
        assert(len(old_children) == len(self.children))
        assert(len(old_beads) == len(self.beads))
349

cmaffeo2's avatar
cmaffeo2 committed
350
    def _generate_beads(self, bead_model, max_basepairs_per_bead, max_nucleotides_per_bead):
351
352
        
        """ Generate beads (positions, types, etcl) and bonds, angles, dihedrals, exclusions """
cmaffeo2's avatar
cmaffeo2 committed
353
        ## TODO: decide whether to remove bead_model argument
354
        ##       (currently unused)
cmaffeo2's avatar
cmaffeo2 committed
355

356
        ## First find points between-which beads must be generated
cmaffeo2's avatar
cmaffeo2 committed
357
358
        conn_locs = self.get_contour_sorted_connections_and_locations()
        locs = [A for c,A,B in conn_locs]
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
        existing_beads = [l.particle for l in locs if l.particle is not None]
        for b in existing_beads:
            assert(b.parent is not None)

        ## Add ends if they don't exist yet
        ## TODO: what about 1 nt segments?
        if len(existing_beads) == 0 or existing_beads[0].get_contour_position(self) > 0:
            b = self._generate_one_bead(0, 0)
            existing_beads = [b] + existing_beads
        if existing_beads[-1].get_contour_position(self) < 1:
            b = self._generate_one_bead(1, 0)
            existing_beads.append(b)
        assert(len(existing_beads) > 1)

        ## Walk through existing_beads, add beads between
        tmp_children = []       # build list of children in nice order
375
        last = None
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
        for I in range(len(existing_beads)-1):
            eb1,eb2 = [existing_beads[i] for i in (I,I+1)]

            # print(" %s working on %d to %d" % (self.name, eb1.position[2], eb2.position[2]))
            e_ds = eb2.get_contour_position(self) - eb1.get_contour_position(self)
            num_beads = self._get_num_beads( e_ds, max_basepairs_per_bead, max_nucleotides_per_bead )
            ds = e_ds / (num_beads+1)
            nts = ds*self.num_nts
            eb1.num_nts += 0.5*nts
            eb2.num_nts += 0.5*nts

            ## Add beads
            if eb1.parent == self:
                tmp_children.append(eb1)

            s0 = eb1.get_contour_position(self)
            if last is not None:
                last.intrahelical_neighbors.append(eb1)
                eb1.intrahelical_neighbors.append(last)
            last = eb1
            for j in range(num_beads):
                s = ds*(j+1) + s0
                b = self._generate_one_bead(s,nts)

                last.intrahelical_neighbors.append(b)
                b.intrahelical_neighbors.append(last)
                last = b
                tmp_children.append(b)

        last.intrahelical_neighbors.append(eb2)
        eb2.intrahelical_neighbors.append(last)

        if eb2.parent == self:
            tmp_children.append(eb2)
        self._rebuild_children(tmp_children)
411
412
413
414
415
416
417
418
419
420
421
422
423

    def _regenerate_beads(self, max_nts_per_bead=4, ):
        ...
    

class DoubleStrandedSegment(Segment):

    """ Class that describes a segment of ssDNA. When built from
    cadnano models, should not span helices """

    def __init__(self, name, num_nts, start_position = np.array((0,0,0)),
                 end_position = None, 
                 segment_model = None,
cmaffeo2's avatar
cmaffeo2 committed
424
425
                 local_twist = False,
                 num_turns = None,
426
                 start_orientation = None):
cmaffeo2's avatar
cmaffeo2 committed
427
428
429
        
        self.helical_rise = 10.44
        self.distance_per_nt = 3.4
430
431
432
433
434
        Segment.__init__(self, name, num_nts, 
                         start_position,
                         end_position, 
                         segment_model)

cmaffeo2's avatar
cmaffeo2 committed
435
436
437
438
439
440
441
442
443
        self.local_twist = local_twist
        if num_turns is None:
            num_turns = float(num_nts) / self.helical_rise
        self.twist_per_nt = float(360 * num_turns) / num_nts

        if start_orientation is None:
            start_orientation = np.array(((1,0,0),(0,1,0),(0,0,1)))
        self.start_orientation = start_orientation

444
445
        self.nicks = []

446
        self.start = self.start5 = Location( self, address=0, type_= "end5" )
447
448
        self.start3 = Location( self, address=0, type_ = "end3" )

449
450
        self.end = self.end5 = Location( self, address=1, type_= "end5" )
        self.end3 = Location( self, address=1, type_ = "end3" )
451

452
453
454
455
456
457
458
459
460
        ## Set up interpolation for azimuthal angles 
        a = np.array([self.start_position,self.end_position]).T
        tck, u = interpolate.splprep( a, u=[0,1], s=0, k=1)
        self.position_spline_params = tck
        
        ## TODO: initialize sensible spline for orientation
        self.quaternion_spline_params = None


461
    ## Convenience methods
462
    ## TODO: add errors if unrealistic connections are made
463
    def connect_start5(self, end3, type_="intrahelical", force_connection=False):
464
465
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
466
467
        self._connect_ends( self.start5, end3, type_, force_connection = force_connection )
    def connect_start3(self, end5, type_="intrahelical", force_connection=False):
468
469
        if isinstance(end5, SingleStrandedSegment):
            end5 = end5.end5
470
471
        self._connect_ends( self.start3, end5, type_, force_connection = force_connection )
    def connect_end3(self, end5, type_="intrahelical", force_connection=False):
472
473
        if isinstance(end5, SingleStrandedSegment):
            end5 = end5.end5
474
475
        self._connect_ends( self.end3, end5, type_, force_connection = force_connection )
    def connect_end5(self, end3, type_="intrahelical", force_connection=False):
476
477
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
478
        self._connect_ends( self.end5, end3, type_, force_connection = force_connection )
479

cmaffeo2's avatar
cmaffeo2 committed
480
    def crossover(self, nt, other, other_nt, strands_fwd=[True,False]):
cmaffeo2's avatar
cmaffeo2 committed
481
482
483
484
485
486
487
        """ Add a crossover between two helices """
        ## Validate other, nt, other_nt
        ##   TODO

        ## Create locations, connections and add to segments
        c = nt/self.num_nts
        assert(c >= 0 and c <= 1)
cmaffeo2's avatar
cmaffeo2 committed
488
        loc = Location( self, address=c, type_ = strands_fwd[0] )
cmaffeo2's avatar
cmaffeo2 committed
489
490
491

        c = other_nt/other.num_nts
        assert(c >= 0 and c <= 1)
cmaffeo2's avatar
cmaffeo2 committed
492
493
        other_loc = Location( other, address=c, type_ = strands_fwd[1] )
        self._connect(other, Connection( loc, other_loc, type_="crossover" ))
cmaffeo2's avatar
cmaffeo2 committed
494

495
    ## Real work
496
    def _connect_ends(self, end1, end2, type_, force_connection):
497
        ## TODO remove self?
498
499
500
501
502
        ## validate the input
        for end in (end1, end2):
            assert( isinstance(end, Location) )
            assert( end.type_ in ("end3","end5") )
        assert( end1.type_ != end2.type_ )
503
        ## Create and add connection
504
        end1.container._connect( end2.container, Connection( end1, end2, type_=type_ ) )
505

506
507
    def _get_num_beads(self, contour, max_basepairs_per_bead, max_nucleotides_per_bead):
        return int(contour*self.num_nts // max_basepairs_per_bead)
cmaffeo2's avatar
cmaffeo2 committed
508

cmaffeo2's avatar
cmaffeo2 committed
509
510
511
512
513
514
515
516
517
518
519
520
521
522
    def get_5prime_locations(self):
        locs = []
        
        ## Add ends of segment if they aren't connected
        cl = self.get_connections_and_locations()
        connlocs = [A for c,A,B in cl]
        if self.start5 not in connlocs:
            locs.append((0,True)) # TODO
        if self.end5 not in connlocs:
            locs.append((1,False)) # TODO
        
        ## TODO Add nicks
        return locs
        
523
524
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
525
        if self.local_twist:
526
            orientation = self.contour_to_orientation(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
527
            opos = pos + orientation.dot( np.array((Segment.orientation_bond.r0,0,0)) )
528
529
530
531
532
533
            o = SegmentParticle( Segment.orientation_particle, opos, nts,
                                 num_nts=nts, parent=self )
            bead = SegmentParticle( Segment.dsDNA_particle, pos, nts,
                                    num_nts=nts, parent=self, 
                                    orientation_bead=o,
                                    contour_position=contour_position )
cmaffeo2's avatar
cmaffeo2 committed
534
535

        else:
536
537
538
539
            bead = SegmentParticle( Segment.dsDNA_particle, pos, nts,
                                    num_nts=nts, parent=self,
                                    contour_position=contour_position )
        self._add_bead(bead)
cmaffeo2's avatar
cmaffeo2 committed
540
        return bead
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
        

class SingleStrandedSegment(Segment):

    """ Class that describes a segment of ssDNA. When built from
    cadnano models, should not span helices """

    def __init__(self, name, num_nts, start_position = np.array((0,0,0)),
                 end_position = None, 
                 segment_model = None):

        self.distance_per_nt = 5
        Segment.__init__(self, name, num_nts, 
                         start_position,
                         end_position, 
                         segment_model)

        self.start = self.end5 = Location( self, address=0, type_= "end5" )
559
        self.end = self.end3 = Location( self, address=1, type_ = "end3" )
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577

    def connect_3end(self, end5, force_connection=False):
        self._connect_end( end5,  _5_to_3 = False, force_connection = force_connection )

    def connect_5end(self, end3, force_connection=False):
        self._connect_end( end3,  _5_to_3 = True, force_connection = force_connection )

    def _connect_end(self, other, _5_to_3, force_connection):
        assert( isinstance(other, Location) )
        if _5_to_3 == True:
            my_end = self.end5
            assert( other.type_ == "end3" )
        else:
            my_end = self.end3
            assert( other.type_ == "end5" )

        self._connect( other.container, Connection( my_end, other, type_="intrahelical" ) )

578
579
    def _get_num_beads(self, contour, max_basepairs_per_bead, max_nucleotides_per_bead):
        return int(contour*self.num_nts // max_nucleotides_per_bead)
cmaffeo2's avatar
cmaffeo2 committed
580

cmaffeo2's avatar
cmaffeo2 committed
581
582
583
584
585
586
587
588
589
590
591
    def get_5prime_locations(self):
        locs = []
        
        ## Add ends of segment if they aren't connected
        cl = self.get_connections_and_locations()
        connLocs = [A for c,A,B in cl]
        if self.start not in connLocs:
            locs.append((0,True)) # TODO

        return locs

592
593
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
594
595
596
597
598
        b = SegmentParticle( Segment.ssDNA_particle, pos, nts,
                             num_nts=nts, parent=self,
                             contour_position=contour_position )
        self._add_bead(b)
        return b
599
600

    
cmaffeo2's avatar
cmaffeo2 committed
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
class StrandInSegment(Group):
    """ Class that holds atomic model, maps to segment """
    
    def __init__(self, segment, start, end, is_fwd):
        """ start/end should be provided expressed as contour_length, is_fwd tuples """
        Group.__init__(self)
        self.num_nts = 0
        self.sequence = []
        self.segment = segment
        self.start = start
        self.end = end
        self.is_fwd = is_fwd

        ## TODO: get sequence (from segment?)        
        nts = np.abs(end-start)*segment.num_nts
        self.num_nts = round(nts)
        assert( np.abs(self.num_nts-nts) < 0.001 )

        print(" Creating {}-nt StrandInSegment in {} from {} to {} {}".format(self.num_nts, segment.name, start, end, is_fwd))

class Strand(Group):
    """ Class that holds atomic model, maps to segments """
    
    def __init__(self):
        Group.__init__(self)
        self.num_nts = 0
        self.children = self.strand_segments = []

    def add_dna(self, segment, start, end, is_fwd):
        """ start/end should be provided expressed as contour_length, is_fwd tuples """
        s = StrandInSegment( segment, start, end, is_fwd )
        self.strand_segments.append( s )
        self.num_nts += s.num_nts

    def generate_atomic_model(self):
        last = None
        resid = 1
        for s in self.strand_segments:
            seg = s.segment
            for c in np.linspace(s.start,s.end,s.num_nts):
                if s.is_fwd:
                    nt = seg._generate_atomic_nucleotide( c, s.is_fwd, "T" ) # TODO add sequence,termini
                else:
                    nt = seg._generate_atomic_nucleotide( c, s.is_fwd, "A" )

                s.children.append(nt)
                ## Join last basepairs
                if last is not None:
                    o3,c3,c4,c2,h3 = [last.atoms_by_name[n] 
                                      for n in ("O3'","C3'","C4'","C2'","H3'")]
                    p,o5,o1,o2,c5 = [nt.atoms_by_name[n] 
                                     for n in ("P","O5'","O1P","O2P","C5'")]
                    self.add_bond( o3, p, None )
                    self.add_angle( c3, o3, p, None )
                    for x in (o5,o1,o2):
                        self.add_angle( o3, p, x, None )
                        self.add_dihedral(c3, o3, p, x, None )
                    for x in (c4,c2,h3):
                        self.add_dihedral(x, c3, o3, p, None )
                    self.add_dihedral(o3, p, o5, c5, None)
                nt.__dict__['resid'] = resid
                resid += 1
                last = nt
664
665

class SegmentModel(ArbdModel):
cmaffeo2's avatar
cmaffeo2 committed
666
667
668
    def __init__(self, segments=[], local_twist=True,
                 max_basepairs_per_bead=7,
                 max_nucleotides_per_bead=4,
669
670
671
                 dimensions=(1000,1000,1000), temperature=291,
                 timestep=50e-6, cutoff=50, 
                 decompPeriod=10000, pairlistDistance=None, 
672
673
674
                 nonbondedResolution=0,DEBUG=0):
        self.DEBUG = DEBUG
        if DEBUG > 0: print("Building ARBD Model")
675
676
677
678
679
        ArbdModel.__init__(self,segments,
                           dimensions, temperature, timestep, cutoff, 
                           decompPeriod, pairlistDistance=None,
                           nonbondedResolution=0)

cmaffeo2's avatar
cmaffeo2 committed
680
681
        # self.max_basepairs_per_bead = max_basepairs_per_bead     # dsDNA
        # self.max_nucleotides_per_bead = max_nucleotides_per_bead # ssDNA
682
683
684
685
        self.children = self.segments = segments

        self._bonded_potential = dict() # cache bonded potentials

686
        self._generate_bead_model( max_basepairs_per_bead, max_nucleotides_per_bead, local_twist)
687
688
689

        self.useNonbondedScheme( nbDnaScheme )

690

cmaffeo2's avatar
cmaffeo2 committed
691
    def get_connections(self,type_=None,exclude=[]):
692
693
694
695
        """ Find all connections in model, without double-counting """
        added=set()
        ret=[]
        for s in self.segments:
cmaffeo2's avatar
cmaffeo2 committed
696
            items = [e for e in s.get_connections_and_locations(type_,exclude=exclude) if e[0] not in added]
697
698
699
            added.update([e[0] for e in items])
            ret.extend( items )
        return ret
700
    
701
    def _recursively_get_beads_within_bonds(self,b1,bonds,done=[]):
702
        ret = []
703
704
705
706
707
708
709
710
711
        done = list(done)
        done.append(b1)
        if bonds == 0:
            return [[]]

        for b2 in b1.intrahelical_neighbors:
            if b2 in done: continue
            for tmp in self._recursively_get_beads_within_bonds(b2, bonds-1, done):
                ret.append( [b2]+tmp )
712
713
        return ret

714
715
716
717
718
    def _get_intrahelical_beads(self,num=2):
        ## TODO: add check that this is not called before adding intrahelical_neighbors in _generate_bead_model

        assert(num >= 2)

719
720
        ret = []
        for s in self.segments:
721
722
723
724
725
            for b1 in s.beads:
                for bead_list in self._recursively_get_beads_within_bonds(b1, num-1):
                    assert(len(bead_list) == num-1)
                    if b1.idx < bead_list[-1].idx: # avoid double-counting
                        ret.append([b1]+bead_list)
726
727
        return ret

728

729
730
    def _get_intrahelical_angle_beads(self):
        return self._get_intrahelical_beads(num=3)
731

cmaffeo2's avatar
cmaffeo2 committed
732
    def _get_potential(self, type_, kSpring, d, max_potential = None):
733
734
735
        key = (type_,kSpring,d)
        if key not in self._bonded_potential:
            if type_ == "bond":
736
                self._bonded_potential[key] = HarmonicBond(kSpring,d, rRange=(0,500), max_potential=max_potential)
737
            elif type_ == "angle":
cmaffeo2's avatar
cmaffeo2 committed
738
739
                self._bonded_potential[key] = HarmonicAngle(kSpring,d, max_potential=max_potential)
                # , resolution = 1, maxForce=0.1)
740
            elif type_ == "dihedral":
cmaffeo2's avatar
cmaffeo2 committed
741
                self._bonded_potential[key] = HarmonicDihedral(kSpring,d, max_potential=max_potential)
742
743
744
745
746
747
748
            else:
                raise Exception("Unhandled potential type '%s'" % type_)
        return self._bonded_potential[key]
    def get_bond_potential(self, kSpring, d):
        return self._get_potential("bond", kSpring, d)
    def get_angle_potential(self, kSpring, d):
        return self._get_potential("angle", kSpring, d)
cmaffeo2's avatar
cmaffeo2 committed
749
750
751
752
    def get_dihedral_potential(self, kSpring, d, max_potential=None):
        while d > 180: d-=360
        while d < -180: d+=360
        return self._get_potential("dihedral", kSpring, d, max_potential)
753
754


cmaffeo2's avatar
cmaffeo2 committed
755
756
757
    def _getParent(self, *beads ):
        ## TODO: test
        if np.all( [b1.parent == b2.parent 
758
                    for b1,b2 in zip(beads[:-1],beads[1:])] ):
cmaffeo2's avatar
cmaffeo2 committed
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
            return beads[0].parent
        else:
            return self

    def _get_twist_spring_constant(self, sep):
        """ sep in nt """
        kT = 0.58622522         # kcal/mol
        twist_persistence_length = 90  # set semi-arbitrarily as there is a large spread in literature
        ## <cos(q)> = exp(-s/Lp) = integrate( cos[x] exp(-A x^2), {x, 0, pi} ) / integrate( exp(-A x^2), {x, 0, pi} )
        ##   Assume A is small
        ## int[B_] :=  Normal[Integrate[ Series[Cos[x] Exp[-B x^2], {B, 0, 1}], {x, 0, \[Pi]}]/
        ##             Integrate[Series[Exp[-B x^2], {B, 0, 1}], {x, 0, \[Pi]}]]

        ## Actually, without assumptions I get fitFun below
        ## From http://www.annualreviews.org/doi/pdf/10.1146/annurev.bb.17.060188.001405
        ##   units "3e-19 erg cm/ 295 k K" "nm" =~ 73
        Lp = twist_persistence_length/0.34 

        fitFun = lambda x: np.real(erf( (4*np.pi*x + 1j)/(2*np.sqrt(x)) )) * np.exp(-1/(4*x)) / erf(2*np.sqrt(x)*np.pi) - np.exp(-sep/Lp)
        k = opt.leastsq( fitFun, x0=np.exp(-sep/Lp) )
        return k[0][0] * 2*kT*0.00030461742
780
781
782
783
784
785

    """ Mapping between different resolution models """
    def _clear_beads(self):
        for s in self.segments:
            s.clear_all()
        self.clear_all(keep_children=True)
786
787
788
789
790
791
792
793
794
795
        assert( len([b for b in self]) == 0 )
        locParticles = []
        # for c,A,B in self.get_connections():
        for s in self.segments:
            for c,A,B in s.get_connections_and_locations():
                for l in (A,B):
                    if l.particle is not None:
                        locParticles.append(A.particle)
        assert( len(locParticles) == 0 )
        assert( len([b for s in self.segments for b in s.beads]) == 0 )
796
797
798
799

    def _update_segment_positions(self, bead_coordinates):
        """ Set new function for each segments functions
        contour_to_position and contour_to_orientation """
cmaffeo2's avatar
cmaffeo2 committed
800

801
        for s in self.segments:
802
803
804
805
806
807
808
809
810
            ## TODO: may need to make spline on helices continuously differentiable
            # beads = s.beads # [b for b in s.beads]
            beads = list(set(s.beads + [A.particle for c,A,B in s.get_connections_and_locations("intrahelical")]))
            contours = [b.get_contour_position(s) for b in beads]

            cb = sorted( zip(contours,beads), key=lambda a:a[0] )
            beads = [b for c,b in cb] 
            contours = [c for c,b in cb] 

811
812
813
814
            ids = [b.idx for b in beads]
            
            """ Get positions """
            positions = bead_coordinates[ids,:].T
815
816
817
818
819
820
821

            ## TODO: determine whether having beads is an issue 
            try:
                tck, u = interpolate.splprep( positions, u=contours, s=0, k=3 )
            except:
                tck, u = interpolate.splprep( positions, u=contours, s=0, k=1 )

822
823
824
            s.position_spline_params = tck

            """ Get twist """
825
826
827
828
829
830
            cb = [e for e in cb if 'orientation_bead' in e[1].__dict__]
            beads = [b for c,b in cb] 
            contours = [c for c,b in cb] 
            ids = [b.idx for b in beads]
            # if 'orientation_bead' in beads[0].__dict__:
            if len(beads) > 3:
831
832
                tangents = s.contour_to_tangent(contours)
                quats = []
cmaffeo2's avatar
cmaffeo2 committed
833
                lastq = None
834
835
836
837
838
839
                for b,t in zip(beads,tangents):
                    o = b.orientation_bead
                    angleVec = o.position - b.position
                    angleVec = angleVec - angleVec.dot(t)*t
                    angleVec = angleVec/np.linalg.norm(angleVec)
                    y = np.cross(t,angleVec)
cmaffeo2's avatar
cmaffeo2 committed
840
841
842
843
844
845
846
847
848
                    assert( np.abs(np.linalg.norm(y) - 1) < 1e-2 )
                    # quats.append( quaternion_from_matrix( np.array([t,y,angleVec])) )
                    q = quaternion_from_matrix( np.array([angleVec,y,t]).T)

                    if lastq is not None:
                        if q.dot(lastq) < 0:
                            q = -q
                    quats.append( q )

849
                quats = np.array(quats)
850
851
852
853
                try:
                    tck, u = interpolate.splprep( quats.T, u=contours, s=0, k=3 )
                except:
                    tck, u = interpolate.splprep( quats.T, u=contours, s=0, k=1 )
854
855
856
857
858
                s.quaternion_spline_params = tck

            ## TODO: set twist

    def _generate_bead_model(self,
859
                             max_basepairs_per_bead = 7,
cmaffeo2's avatar
cmaffeo2 committed
860
                             max_nucleotides_per_bead = 4,
861
                             local_twist=False):
862

863
        segments = self.segments
864
865
866
        for s in segments:
            s.local_twist = local_twist
            
867
        """ Generate beads at intrahelical junctions """
868
869
870
871
872
873
874
875
876
877
878
879
880
881
        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"):
            s1,s2 = [l.container for l in (A,B)]
            if isinstance(s1,DoubleStrandedSegment) and isinstance(s2,DoubleStrandedSegment):
                assert( A.particle is None )
                assert( B.particle is None )

                ## TODO: offload the work here to s1
                a1,a2 = [l.address   for l in (A,B)]
                a1,a2 = [a - (0.5/s.num_nts) if a == 0 else a + (0.5/s.num_nts) for a,s in zip((a1,a2),(s1,s2))]

                b = s1._generate_one_bead(a1,0)
                A.particle = B.particle = b
882

883
            else:
cmaffeo2's avatar
cmaffeo2 committed
884
                ## TODO fix this for ssDNA vs dsDNA (maybe a1/a2 should be calculated differently)
885
886
                a1,a2 = [l.address   for l in (A,B)]
                a1,a2 = [a - (0.5/s.num_nts) if a == 0 else a + (0.5/s.num_nts) for a,s in zip((a1,a2),(s1,s2))]
cmaffeo2's avatar
cmaffeo2 committed
887
888
889
890
                if isinstance(s2,DoubleStrandedSegment):
                    b = s2._generate_one_bead(a2,0)
                else:
                    b = s1._generate_one_bead(a1,0)
891
892
                A.particle = B.particle = b

893
        """ Generate beads at other junctions """
cmaffeo2's avatar
cmaffeo2 committed
894
895
896
897
898
899
900
        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:
                continue
            assert( A.particle is None )
            assert( B.particle is None )

901
            ## TODO: offload the work here to s1/s2 (?)
cmaffeo2's avatar
cmaffeo2 committed
902
            a1,a2 = [l.address   for l in (A,B)]
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918

            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)
cmaffeo2's avatar
cmaffeo2 committed
919

920
921
922
923
924
925
926
        """ Some tests """
        for c,A,B in self.get_connections("intrahelical"):
            for l in (A,B):
                if l.particle is None: continue
                assert( l.particle.parent is not None )

        """ Generate beads in between """
927
        if self.DEBUG: print("Generating beads")
928
        for s in segments:
cmaffeo2's avatar
cmaffeo2 committed
929
            s._generate_beads( self, max_basepairs_per_bead, max_nucleotides_per_bead )
930

931
932
933
934
935
936
937
938
939
940
941
942
943
944
        # """ Combine beads at junctions as needed """
        # for c,A,B in self.get_connections():
        #    ...

        """ Add intrahelical neighbors at connections """
        for c,A,B in self.get_connections("intrahelical"):
            b1,b2 = [l.particle for l in (A,B)]
            if b1 is b2:
                ## already handled by Segment._generate_beads
                continue
            else:
                for b in (b1,b2): assert( b is not None )
                b1.intrahelical_neighbors.append(b2)
                b2.intrahelical_neighbors.append(b1)
945
946

        """ Reassign bead types """
947
        if self.DEBUG: print("Assigning bead types")
948
949
950
        beadtype_s = dict()
        for segment in segments:
            for b in segment:
951
                b.num_nts = np.around( b.num_nts, decimals=1 )
952
953
954
955
956
957
958
959
960
                key = (b.type_.name[0].upper(), b.num_nts)
                if key in beadtype_s:
                    b.type_ = beadtype_s[key]
                else:
                    t = deepcopy(b.type_)
                    if key[0] == "D":
                        t.__dict__["nts"] = b.num_nts*2
                    elif key[0] == "S":
                        t.__dict__["nts"] = b.num_nts
cmaffeo2's avatar
cmaffeo2 committed
961
962
                    elif key[0] == "O":
                        t.__dict__["nts"] = b.num_nts
963
964
                    else:
                        raise Exception("TODO")
cmaffeo2's avatar
cmaffeo2 committed
965
                    # print(t.nts)
966
                    t.name = t.name + "%03d" % (10*t.nts)
967
968
                    beadtype_s[key] = b.type_ = t

969
970
971
        """ Update bead indices """
        self._countParticleTypes() # probably not needed here
        self._updateParticleOrder()
972

973
        """ Add intrahelical bond potentials """
974
        if self.DEBUG: print("Adding intrahelical bond potentials")
cmaffeo2's avatar
cmaffeo2 committed
975
        dists = dict()          # intrahelical distances built for later use
976
977
978
        intra_beads = self._get_intrahelical_beads() 
        if self.DEBUG: print("  Adding %d bonds" % len(intra_beads))
        for b1,b2 in intra_beads:
cmaffeo2's avatar
cmaffeo2 committed
979
            parent = self._getParent(b1,b2)
980
981
982

            ## TODO: could be sligtly smarter about sep
            sep = 0.5*(b1.num_nts+b2.num_nts)
cmaffeo2's avatar
cmaffeo2 committed
983

984
            conversion = 0.014393265 # units "pN/AA" kcal_mol/AA^2
985
            if b1.type_.name[0] == "D" and b2.type_.name[0] == "D":
986
                elastic_modulus = 1000 # pN http://markolab.bmbcb.northwestern.edu/marko/Cocco.CRP.02.pdf
987
                d = 3.4*sep
988
                k = conversion*elastic_modulus/d
989
            else:
990
991
                ## TODO: get better numbers our ssDNA model
                elastic_modulus = 800 # pN http://markolab.bmbcb.northwestern.edu/marko/Cocco.CRP.02.pdf
992
                d = 5*sep
993
                k = conversion*elastic_modulus/d
994
                # print(sep,d,k)
995
996
              
            if b1 not in dists:
cmaffeo2's avatar
cmaffeo2 committed
997
                dists[b1] = dict()
998
            if b2 not in dists:
cmaffeo2's avatar
cmaffeo2 committed
999
1000
1001
1002
1003
                dists[b2] = dict()
            # dists[b1].append([b2,sep])
            # dists[b2].append([b1,sep])
            dists[b1][b2] = sep
            dists[b2][b1] = sep
1004
1005

            # dists[[b1,b2]] = dists[[b2,b1]] = sep
1006
1007
            bond = self.get_bond_potential(k,d)
            parent.add_bond( b1, b2, bond, exclude=True )
1008
1009

        """ Add intrahelical angle potentials """
1010
        if self.DEBUG: print("Adding intrahelical angle potentials")
1011
        for b1,b2,b3 in self._get_intrahelical_angle_beads():
1012
1013
            ## TODO: could be slightly smarter about sep
            sep = 0.5*b1.num_nts+b2.num_nts+0.5*b3.num_nts
cmaffeo2's avatar
cmaffeo2 committed
1014
            parent = self._getParent(b1,b2,b3)
1015

1016
1017
1018
1019
1020
1021
            kT = 0.58622522         # kcal/mol
            if b1.type_.name[0] == "D" and b2.type_.name[0] == "D" and b3.type_.name[0] == "D":
                ## <cos(q)> = exp(-s/Lp) = integrate( x^4 exp(-A x^2) / 2, {x, 0, pi} ) / integrate( x^2 exp(-A x^2), {x, 0, pi} )
                ## <cos(q)> ~ 1 - 3/4A                                                                                            
                ## where A = k_spring / (2 kT)                                                                                    
                k = 1.5 * kT * (1.0 / (1-np.exp(-float(sep)/147))) * 0.00030461742; # kcal_mol/degree^2
cmaffeo2's avatar
cmaffeo2 committed
1022
1023
                if local_twist:
                    k *= 0.5    # halve because orientation beads have similar springs
1024
1025
            else:
                ## TODO: get correct number from ssDNA model
cmaffeo2's avatar
cmaffeo2 committed
1026
                k = 1.5 * kT * (1.0 / (1-np.exp(-float(sep)/3))) * 0.00030461742; # kcal_mol/degree^2
1027

1028
1029
1030
1031
            angle = self.get_angle_potential(k,180)
            parent.add_angle( b1, b2, b3, angle )

        """ Add intrahelical exclusions """
1032
        if self.DEBUG: print("Adding intrahelical exclusions")
1033
1034
1035
        beads = dists.keys()
        def _recursively_get_beads_within(b1,d,done=[]):
            ret = []
cmaffeo2's avatar
cmaffeo2 committed
1036
            for b2,sep in dists[b1].items():
1037
1038
1039
1040
1041
1042
1043
                if b2 in done: continue
                if sep < d:
                    ret.append( b2 )
                    done.append( b2 )
                    tmp = _recursively_get_beads_within(b2, d-sep, done)
                    if len(tmp) > 0: ret.extend(tmp)
            return ret
1044

1045
1046
1047
        exclusions = set()
        for b1 in beads:
            exclusions.update( [(b1,b) for b in _recursively_get_beads_within(b1, 20, done=[b1])] )
1048
1049

        if self.DEBUG: print("Adding %d exclusions" % len(exclusions))
1050
        for b1,b2 in exclusions:
cmaffeo2's avatar
cmaffeo2 committed
1051
            parent = self._getParent(b1,b2)
1052
            parent.add_exclusion( b1, b2 )
cmaffeo2's avatar
cmaffeo2 committed
1053
1054
1055

        """ Twist potentials """
        if local_twist:
1056
            if self.DEBUG: print("Adding twist potentials")
cmaffeo2's avatar
cmaffeo2 committed
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066

            ## TODO: decide whether to add bond here
            # """Add bonds between orientation bead and parent"""
            # for s in self.segments:
            #     for b,o in zip(s.children[::2],s.children[1::2]):
            #         s.add_bond(
                    

            for b1 in beads:
                if "orientation_bead" not in b1.__dict__: continue
cmaffeo2's avatar
cmaffeo2 committed
1067
                for b2,sep in dists[b1].items():
cmaffeo2's avatar
cmaffeo2 committed
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
                    if "orientation_bead" not in b2.__dict__: continue

                    p1,p2 = [b.parent for b in (b1,b2)]
                    o1,o2 = [b.orientation_bead for b in (b1,b2)]

                    parent = self._getParent( b1, b2 )
                    

                    """ Add heuristic 90 degree potential to keep orientation bead orthogonal """
                    k = (1.0/2) * 1.5 * kT * (1.0 / (1-np.exp(-float(sep)/147))) * 0.00030461742; # kcal_mol/degree^2
                    pot = self.get_angle_potential(k,90)
                    parent.add_angle(o1,b1,b2, pot)
                    parent.add_angle(b1,b2,o2, pot)
                                        
                    ## TODO: improve this
                    twist_per_nt = 0.5 * (p1.twist_per_nt + p2.twist_per_nt)
                    angle = sep*twist_per_nt
                    if angle > 360 or angle < -360:
                        raise Exception("The twist between beads is too large")
                        
                    k = self._get_twist_spring_constant(sep)
                    pot = self.get_dihedral_potential(k,angle,max_potential=1)
                    parent.add_dihedral(o1,b1,b2,o2, pot)

1092
        """ Add connection potentials """
1093
        for c,A,B in self.get_connections("terminal_crossover"):
cmaffeo2's avatar
cmaffeo2 committed
1094
            ## TODO: use a better description here
1095
1096
1097
            b1,b2 = [loc.particle for loc in (c.A,c.B)]
            pot = self.get_bond_potential(4,18.5)
            self.add_bond(b1,b2, pot)
1098

cmaffeo2's avatar
cmaffeo2 committed
1099
        for c,A,B in self.get_connections("crossover"):
cmaffeo2's avatar
cmaffeo2 committed
1100
            ## TODO: avoid double-counting for double-crossovers
cmaffeo2's avatar
cmaffeo2 committed
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
            b1,b2 = [loc.particle for loc in (c.A,c.B)]
            pot = self.get_bond_potential(4,18.5)
            self.add_bond(b1,b2, pot)

            if not local_twist:
                ## TODO
                ...
            else:
                k = (1.0/2) * 1.5 * kT * (1.0 / (1-np.exp(-float(1)/147))) * 0.00030461742; # kcal_mol/degree^2
                if 'orientation_bead' in b1.__dict__:
                    t0 = 90 + 60
cmaffeo2's avatar
cmaffeo2 committed
1112
                    if A.type_: t0 -= 120
cmaffeo2's avatar
cmaffeo2 committed
1113
1114
1115
1116
1117
                    o = b1.orientation_bead
                    pot = self.get_angle_potential(k,t0)
                    self.add_angle( o,b1,b2, pot )
                else:
                    t0 = 90 + 60
cmaffeo2's avatar
cmaffeo2 committed
1118
                    if B.type_: t0 -= 120
cmaffeo2's avatar
cmaffeo2 committed
1119
1120
1121
1122
1123
                    o = b2.orientation_bead
                    pot = self.get_angle_potential(k,t0)
                    self.add_angle( b1,b2,o, pot )


cmaffeo2's avatar
cmaffeo2 committed
1124
1125
                ## TODO: fix the following to ensure that neighbors are indeed above and below
                ### preferably with a function call to b
cmaffeo2's avatar
cmaffeo2 committed
1126
1127
1128
                u1,u2 = [b.intrahelical_neighbors[-1] for b in (b1,b2)] # TODO: fix this
                d1,d2 = [b.intrahelical_neighbors[0] for b in (b1,b2)] # TODO: fix this

cmaffeo2's avatar
cmaffeo2 committed
1129
                k_fn = lambda sep: (1.0/2) * 1.5 * kT * (1.0 / (1-np.exp(-float(sep)/147))) * 0.00030461742; # kcal_mol/degree^2
cmaffeo2's avatar
cmaffeo2 committed
1130
1131
1132
1133
                t0 = 90
                if 'orientation_bead' in b1.__dict__:
                    o1 = b1.orientation_bead
                    if u2 is not None:
cmaffeo2's avatar
cmaffeo2 committed
1134
1135
                        k = k_fn( dists[b2][u2] )
                        pot = self.get_dihedral_potential(k,t0)
cmaffeo2's avatar
cmaffeo2 committed
1136
1137
                        self.add_dihedral( o1,b1,b2,u2, pot )
                    elif d2 is not None:
cmaffeo2's avatar
cmaffeo2 committed
1138
1139
1140
                        k = k_fn( dists[b2][d2] )
                        pot = self.get_dihedral_potential(k,t0)
                        self.add_dihedral( o1,b1,b2,d2, pot )
cmaffeo2's avatar
cmaffeo2 committed
1141
1142
1143
                if 'orientation_bead' in b2.__dict__:
                    o2 = b2.orientation_bead
                    if u1 is not None:
cmaffeo2's avatar
cmaffeo2 committed
1144
1145
                        k = k_fn( dists[b1][u1] )
                        pot = self.get_dihedral_potential(k,t0)
cmaffeo2's avatar
cmaffeo2 committed
1146
1147
                        self.add_dihedral( o2,b2,b1,u1, pot )
                    elif d1 is not None:
cmaffeo2's avatar
cmaffeo2 committed
1148
1149
                        k = k_fn( dists[b1][d1] )
                        pot = self.get_dihedral_potential(k,t0)
cmaffeo2's avatar
cmaffeo2 committed
1150
1151
1152
1153
1154
                        self.add_dihedral( o2,b2,b1,d1, pot )

                if 'orientation_bead' in b1.__dict__ and 'orientation_bead' in b2.__dict__:
                    if u1 is not None and u2 is not None:
                        t0 = 0
cmaffeo2's avatar
cmaffeo2 committed
1155
                        k = k_fn( 0.5*(dists[b1][u1]+dists[b2][u2]) )
cmaffeo2's avatar
cmaffeo2 committed
1156
1157
1158
1159
                        pot = self.get_dihedral_potential(k,t0)
                        self.add_dihedral( u1,b1,b2,u2,  pot )
                    elif d1 is not None and d2 is not None:
                        t0 = 0
cmaffeo2's avatar
cmaffeo2 committed
1160
                        k = k_fn( 0.5*(dists[b1][d1]+dists[b2][d2]) )
cmaffeo2's avatar
cmaffeo2 committed
1161
1162
1163
1164
                        pot = self.get_dihedral_potential(k,t0)
                        self.add_dihedral( d1,b1,b2,d2, pot )
                    elif d1 is not None and u2 is not None:
                        t0 = 180
cmaffeo2's avatar
cmaffeo2 committed
1165
                        k = k_fn( 0.5*(dists[b1][d1]+dists[b2][u2]) )
cmaffeo2's avatar
cmaffeo2 committed
1166
1167
1168
1169
                        pot = self.get_dihedral_potential(k,t0)
                        self.add_dihedral( d1,b1,b2,u2, pot )
                    elif u1 is not None and d2 is not None:
                        t0 = 180
cmaffeo2's avatar
cmaffeo2 committed
1170
                        k = k_fn( 0.5*(dists[b1][u1]+dists[b2][d2]) )
cmaffeo2's avatar
cmaffeo2 committed
1171
1172
                        pot = self.get_dihedral_potential(k,t0)
                        self.add_dihedral( u1,b1,b2,d2, pot )
1173
    
cmaffeo2's avatar
cmaffeo2 committed
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
    def _generate_atomic_model(self):
        self.children = self.strands = strands = []

        """ Build strands from connectivity of helices """
        ## Find all 5-prime ends
        ## TODO sequence

        def _recursively_build_strand(strand, segment, pos, is_fwd):
            s,seg = [strand, segment]
            end_pos, next_seg, next_pos, next_dir = seg.get_end_of_strand(pos, is_fwd)
            s.add_dna(seg, pos, end_pos, is_fwd)
            if next_seg is not None:
                # print("  next_dir: {}".format(next_dir))
                _recursively_build_strand(s, next_seg, next_pos, next_dir)

        for seg in self.segments:
            locs = seg.get_5prime_locations()
            for pos, is_fwd in locs:
                s = Strand()
                _recursively_build_strand(s, seg, pos, is_fwd)
                print("{} {}".format(seg.name,s.num_nts))
                strands.append(s)
        
        self.children = self.strands = sorted(strands, key=lambda s:s.num_nts)[::-1] # or something
        for s in strands:
            s.generate_atomic_model()