segmentmodel.py 48.9 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
19
20
21
22
 - map to atomic representation
 - remove performance bottlenecks
 - test for large systems
 - assign sequence
23
 - document
24
25
26
27
28
"""

class Location():
    """ Site for connection within an object """
    def __init__(self, container, address, type_):
29
        ## TODO: remove cyclic references(?)
30
        self.container = container
31
        self.address = address  # represents position along contour length in segments
cmaffeo2's avatar
cmaffeo2 committed
32
        # assert( type_ in ("end3","end5") ) # TODO remove or make conditional
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
        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
51
    def get_connections_and_locations(self, type_=None, exclude=[]):
52
53
54
55
        """ 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
56
            if type_ is None or c.type_ == type_ and type_ not in exclude:
57
                if   c.A.container is self:
58
                    ret.append( [c, c.A, c.B] )
59
                elif c.B.container is self:
60
61
62
63
64
                    ret.append( [c, c.B, c.A] )
                else:
                    raise Exception("Object contains connection that fails to refer to object")
        return ret

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

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
90
            ## ERROR
91
92
93
94
95
96
97
98
99
100
            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))
101
102
103
104
105
106
107
108
109
110
111
112
       
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
113
114
115
116
117
    orientation_particle = ParticleType("O",
                                        diffusivity = 100,
                                        mass = 300,
                                        radius = 1,
                                    )
118

cmaffeo2's avatar
cmaffeo2 committed
119
    # orientation_bond = HarmonicBond(10,2)
120
    orientation_bond = HarmonicBond(30,1.5, rRange = (0,500) )
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

    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
136
137
138
139
140
        self.start_orientation = None
        self.twist_per_nt = 0

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

141
142
143
        self._bead_model_generation = 0    # TODO: remove?
        self.segment_model = segment_model # TODO: remove?

cmaffeo2's avatar
cmaffeo2 committed
144
        self.num_nts = int(num_nts)
145
146
147
148
149
        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

150
151
152
153
154
        ## 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

155
156
157
158
159
    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
160
161
162
163
164
165
166
167

    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 )
168
169
        t = (t / np.linalg.norm(t,axis=0))
        return t.T
170
171
172
173
174
175
176
        

    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)
177
                orientation = rotationAboutAxis( axis, s*self.twist_per_nt*self.num_nts, normalizeAxis=True )
178
179
180
181
182
                ## 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
183
                q = q/np.linalg.norm(q)
184
                orientation = quaternion_to_matrix(q)
cmaffeo2's avatar
cmaffeo2 committed
185
                
186
187
188
189
        else:
            orientation = None
        return orientation

cmaffeo2's avatar
cmaffeo2 committed
190
191
192
193
    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)
194

cmaffeo2's avatar
cmaffeo2 committed
195
196
197
    def _get_num_beads(self, max_basepairs_per_bead, max_nucleotides_per_bead ):
        raise NotImplementedError

198
    def _generate_one_bead(self, contour_position, nts):
199
200
        raise NotImplementedError

cmaffeo2's avatar
cmaffeo2 committed
201
202
203
204
205
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
    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
        
295
296
297
298
299
300
301
    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]
302
303
304

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

312
313
314
    def _add_bead(self,b,set_contour=False):
        if set_contour:
            b.contour_position = b.get_contour_position(self)
315
        
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
        # 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))
344

cmaffeo2's avatar
cmaffeo2 committed
345
    def _generate_beads(self, bead_model, max_basepairs_per_bead, max_nucleotides_per_bead):
346
347
        
        """ Generate beads (positions, types, etcl) and bonds, angles, dihedrals, exclusions """
cmaffeo2's avatar
cmaffeo2 committed
348
        ## TODO: decide whether to remove bead_model argument
349
        ##       (currently unused)
cmaffeo2's avatar
cmaffeo2 committed
350

351
        ## First find points between-which beads must be generated
cmaffeo2's avatar
cmaffeo2 committed
352
353
        conn_locs = self.get_contour_sorted_connections_and_locations()
        locs = [A for c,A,B in conn_locs]
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
        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
370
        last = None
371
372
373
374
375
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
        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)
406
407
408
409
410
411
412
413
414
415
416
417
418

    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
419
420
                 local_twist = False,
                 num_turns = None,
421
                 start_orientation = None):
cmaffeo2's avatar
cmaffeo2 committed
422
423
424
        
        self.helical_rise = 10.44
        self.distance_per_nt = 3.4
425
426
427
428
429
        Segment.__init__(self, name, num_nts, 
                         start_position,
                         end_position, 
                         segment_model)

cmaffeo2's avatar
cmaffeo2 committed
430
431
432
433
434
435
436
437
438
        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

439
440
        self.nicks = []

441
        self.start = self.start5 = Location( self, address=0, type_= "end5" )
442
443
        self.start3 = Location( self, address=0, type_ = "end3" )

444
445
        self.end = self.end5 = Location( self, address=1, type_= "end5" )
        self.end3 = Location( self, address=1, type_ = "end3" )
446

447
448
449
450
451
452
453
454
455
        ## 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


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

cmaffeo2's avatar
cmaffeo2 committed
475
    def crossover(self, nt, other, other_nt, strands_fwd=[True,False]):
cmaffeo2's avatar
cmaffeo2 committed
476
477
478
479
480
481
482
        """ 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
483
        loc = Location( self, address=c, type_ = strands_fwd[0] )
cmaffeo2's avatar
cmaffeo2 committed
484
485
486

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

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

501
502
    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
503

cmaffeo2's avatar
cmaffeo2 committed
504
505
506
507
508
509
510
511
512
513
514
515
516
517
    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
        
518
519
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
520
        if self.local_twist:
521
            orientation = self.contour_to_orientation(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
522
            opos = pos + orientation.dot( np.array((Segment.orientation_bond.r0,0,0)) )
523
524
525
526
527
528
            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
529
530

        else:
531
532
533
534
            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
535
        return bead
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
        

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" )
554
        self.end = self.end3 = Location( self, address=1, type_ = "end3" )
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572

    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" ) )

573
574
    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
575

cmaffeo2's avatar
cmaffeo2 committed
576
577
578
579
580
581
582
583
584
585
586
    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

587
588
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
589
590
591
592
593
        b = SegmentParticle( Segment.ssDNA_particle, pos, nts,
                             num_nts=nts, parent=self,
                             contour_position=contour_position )
        self._add_bead(b)
        return b
594
595

    
cmaffeo2's avatar
cmaffeo2 committed
596
597
598
599
600
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
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
659
660

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

cmaffeo2's avatar
cmaffeo2 committed
675
676
        # self.max_basepairs_per_bead = max_basepairs_per_bead     # dsDNA
        # self.max_nucleotides_per_bead = max_nucleotides_per_bead # ssDNA
677
678
679
680
        self.children = self.segments = segments

        self._bonded_potential = dict() # cache bonded potentials

681
        self._generate_bead_model( max_basepairs_per_bead, max_nucleotides_per_bead, local_twist)
682
683
684

        self.useNonbondedScheme( nbDnaScheme )

685

cmaffeo2's avatar
cmaffeo2 committed
686
    def get_connections(self,type_=None,exclude=[]):
687
688
689
690
        """ Find all connections in model, without double-counting """
        added=set()
        ret=[]
        for s in self.segments:
cmaffeo2's avatar
cmaffeo2 committed
691
            items = [e for e in s.get_connections_and_locations(type_,exclude=exclude) if e[0] not in added]
692
693
694
            added.update([e[0] for e in items])
            ret.extend( items )
        return ret
695
    
696
    def _recursively_get_beads_within_bonds(self,b1,bonds,done=[]):
697
        ret = []
698
699
700
701
702
703
704
705
706
        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 )
707
708
        return ret

709
710
711
712
713
    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)

714
715
        ret = []
        for s in self.segments:
716
717
718
719
720
            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)
721
722
        return ret

723

724
725
    def _get_intrahelical_angle_beads(self):
        return self._get_intrahelical_beads(num=3)
726

cmaffeo2's avatar
cmaffeo2 committed
727
    def _get_potential(self, type_, kSpring, d, max_potential = None):
728
729
730
        key = (type_,kSpring,d)
        if key not in self._bonded_potential:
            if type_ == "bond":
731
                self._bonded_potential[key] = HarmonicBond(kSpring,d, rRange=(0,500), max_potential=max_potential)
732
            elif type_ == "angle":
cmaffeo2's avatar
cmaffeo2 committed
733
734
                self._bonded_potential[key] = HarmonicAngle(kSpring,d, max_potential=max_potential)
                # , resolution = 1, maxForce=0.1)
735
            elif type_ == "dihedral":
cmaffeo2's avatar
cmaffeo2 committed
736
                self._bonded_potential[key] = HarmonicDihedral(kSpring,d, max_potential=max_potential)
737
738
739
740
741
742
743
            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
744
745
746
747
    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)
748
749


cmaffeo2's avatar
cmaffeo2 committed
750
751
752
    def _getParent(self, *beads ):
        ## TODO: test
        if np.all( [b1.parent == b2.parent 
753
                    for b1,b2 in zip(beads[:-1],beads[1:])] ):
cmaffeo2's avatar
cmaffeo2 committed
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
            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
775
776
777
778
779
780

    """ Mapping between different resolution models """
    def _clear_beads(self):
        for s in self.segments:
            s.clear_all()
        self.clear_all(keep_children=True)
781
782
783
784
785
786
787
788
789
790
        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 )
791
792
793
794

    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
795

796
        for s in self.segments:
797
798
799
800
801
802
803
804
805
            ## 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] 

806
807
808
809
            ids = [b.idx for b in beads]
            
            """ Get positions """
            positions = bead_coordinates[ids,:].T
810
811
812
813
814
815
816

            ## 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 )

817
818
819
            s.position_spline_params = tck

            """ Get twist """
820
821
822
823
824
825
            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:
826
827
                tangents = s.contour_to_tangent(contours)
                quats = []
cmaffeo2's avatar
cmaffeo2 committed
828
                lastq = None
829
830
831
832
833
834
                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
835
836
837
838
839
840
841
842
843
                    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 )

844
                quats = np.array(quats)
845
846
847
848
                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 )
849
850
851
852
853
                s.quaternion_spline_params = tck

            ## TODO: set twist

    def _generate_bead_model(self,
854
                             max_basepairs_per_bead = 7,
cmaffeo2's avatar
cmaffeo2 committed
855
                             max_nucleotides_per_bead = 4,
856
                             local_twist=False):
857

858
        segments = self.segments
859
860
861
        for s in segments:
            s.local_twist = local_twist
            
862
        """ Generate beads at intrahelical junctions """
863
864
865
866
867
868
869
870
871
872
873
874
875
876
        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
877

878
            else:
cmaffeo2's avatar
cmaffeo2 committed
879
                ## TODO fix this for ssDNA vs dsDNA (maybe a1/a2 should be calculated differently)
880
881
                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
882
883
884
885
                if isinstance(s2,DoubleStrandedSegment):
                    b = s2._generate_one_bead(a2,0)
                else:
                    b = s1._generate_one_bead(a1,0)
886
887
                A.particle = B.particle = b

888
        """ Generate beads at other junctions """
cmaffeo2's avatar
cmaffeo2 committed
889
890
891
892
893
894
895
        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 )

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

            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
914

915
916
917
918
919
920
921
        """ 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 """
922
        if self.DEBUG: print("Generating beads")
923
        for s in segments:
cmaffeo2's avatar
cmaffeo2 committed
924
            s._generate_beads( self, max_basepairs_per_bead, max_nucleotides_per_bead )
925

926
927
928
929
930
931
932
933
934
935
936
937
938
939
        # """ 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)
940
941

        """ Reassign bead types """
942
        if self.DEBUG: print("Assigning bead types")
943
944
945
        beadtype_s = dict()
        for segment in segments:
            for b in segment:
946
                b.num_nts = np.around( b.num_nts, decimals=1 )
947
948
949
950
951
952
953
954
955
                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
956
957
                    elif key[0] == "O":
                        t.__dict__["nts"] = b.num_nts
958
959
                    else:
                        raise Exception("TODO")
cmaffeo2's avatar
cmaffeo2 committed
960
                    # print(t.nts)
961
                    t.name = t.name + "%03d" % (10*t.nts)
962
963
                    beadtype_s[key] = b.type_ = t

964
965
966
        """ Update bead indices """
        self._countParticleTypes() # probably not needed here
        self._updateParticleOrder()
967

968
        """ Add intrahelical bond potentials """
969
        if self.DEBUG: print("Adding intrahelical bond potentials")
cmaffeo2's avatar
cmaffeo2 committed
970
        dists = dict()          # intrahelical distances built for later use
971
972
973
        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
974
            parent = self._getParent(b1,b2)
975
976
977

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

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

            # dists[[b1,b2]] = dists[[b2,b1]] = sep
For faster browsing, not all history is shown. View entire blame