segmentmodel.py 91.5 KB
Newer Older
cmaffeo2's avatar
cmaffeo2 committed
1
import pdb
2
import numpy as np
3
import random
4
from arbdmodel import PointParticle, ParticleType, Group, ArbdModel
5
from coords import rotationAboutAxis, quaternion_from_matrix, quaternion_to_matrix
6
7
8
9
from nonbonded import *
from copy import copy, deepcopy
from nbPot import nbDnaScheme

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

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

17
# import pdb
18
"""
cmaffeo2's avatar
cmaffeo2 committed
19
TODO:
cmaffeo2's avatar
cmaffeo2 committed
20
 + fix handling of crossovers for atomic representation
cmaffeo2's avatar
cmaffeo2 committed
21
 + map to atomic representation
22
23
    + add nicks
    - transform ssDNA nucleotides 
cmaffeo2's avatar
cmaffeo2 committed
24
25
    - shrink ssDNA
    - shrink dsDNA backbone
26
    + make orientation continuous
cmaffeo2's avatar
cmaffeo2 committed
27
    - sequence
28
    - handle circular dna
29
 + ensure crossover bead potentials aren't applied twice 
30
 + remove performance bottlenecks
31
32
 - test for large systems
 - assign sequence
33
 - ENM
34
35
 - rework Location class 
 - remove recursive calls
36
 - document
37
 - add unit test of helices connected to themselves
38
39
"""

cmaffeo2's avatar
cmaffeo2 committed
40
41
42
class ParticleNotConnectedError(Exception):
    pass

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

56
57
        self.prev_in_strand = None
        self.next_in_strand = None
58
59
        
        self.combine = None     # some locations might be combined in bead model 
60
61
62
63
64
65
66

    def get_connected_location(self):
        if self.connection is None:
            return None
        else:
            return self.connection.other(self)

67
    def set_connection(self, connection, is_3prime_side_of_connection):
68
        self.connection = connection # TODO weakref? 
69
        self.is_3prime_side_of_connection = is_3prime_side_of_connection
70

71
72
73
74
75
76
77
78
79
80
81
82
    def get_nt_pos(self):
        try:
            pos = self.container.contour_to_nt_pos(self.address, round_nt=True)
        except:
            if self.address == 0:
                pos = 0
            elif self.address == 1:
                pos = self.container.num_nts-1
            else:
                raise
        return pos

83
84
85
86
87
88
89
    def __repr__(self):
        if self.on_fwd_strand:
            on_fwd = "on_fwd_strand"
        else:
            on_fwd = "on_rev_strand"
        return "<Location {}.{}[{:.2f},{:d}]>".format( self.container.name, self.type_, self.address, self.on_fwd_strand)
        
90
91
92
93
94
95
96
97
98
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_
        
99
100
101
102
103
104
105
    def other(self, location):
        if location is self.A:
            return self.B
        elif location is self.B:
            return self.A
        else:
            raise Exception("OutOfBoundsError")
cmaffeo2's avatar
cmaffeo2 committed
106
107
108
109
110

    def __repr__(self):
        return "<Connection {}--{}--{}]>".format( self.A, self.type_, self.B )
        

111
        
112
113
114
# class ConnectableElement(Transformable):
class ConnectableElement():
    """ Abstract base class """
115
116
117
118
    ## TODO: eliminate mutable default arguments
    def __init__(self, connection_locations=[], connections=[]):
        ## TODO decide on names
        self.locations = self.connection_locations = connection_locations
119
120
        self.connections = connections

121
122
123
124
125
126
127
128
129
    def get_locations(self, type_=None, exclude=[]):
        locs = [l for l in self.connection_locations if (type_ is None or l.type_ == type_) and l.type_ not in exclude]
        counter = dict()
        for l in locs:
            if l in counter:
                counter[l] += 1
            else:
                counter[l] = 1
        assert( np.all( [counter[l] == 1 for l in locs] ) )
130
131
132
133
134
135
136
137
138
139
140
141
        return locs

    def get_location_at(self, address, on_fwd_strand=True, new_type="crossover"):
        loc = None
        if (self.num_nts == 1):
            # import pdb
            # pdb.set_trace()
            ## Assumes that intrahelical connections have been made before crossovers
            for l in self.locations:
                if l.on_fwd_strand == on_fwd_strand and l.connection is None:
                    assert(loc is None)
                    loc = l
cmaffeo2's avatar
cmaffeo2 committed
142
            # assert( loc is not None )
143
144
145
146
147
148
149
150
        else:
            for l in self.locations:
                if l.address == address and l.on_fwd_strand == on_fwd_strand:
                    assert(loc is None)
                    loc = l
        if loc is None:
            loc = Location( self, address=address, type_=new_type, on_fwd_strand=on_fwd_strand )
        return loc
151
152

    def get_connections_and_locations(self, connection_type=None, exclude=[]):
153
154
        """ Returns a list with each entry of the form:
            connection, location_in_self, location_in_other """
155
        type_ = connection_type
156
157
        ret = []
        for c in self.connections:
158
            if (type_ is None or c.type_ == type_) and c.type_ not in exclude:
159
                if   c.A.container is self:
160
                    ret.append( [c, c.A, c.B] )
161
                elif c.B.container is self:
162
163
                    ret.append( [c, c.B, c.A] )
                else:
164
165
                    import pdb
                    pdb.set_trace()
166
167
168
                    raise Exception("Object contains connection that fails to refer to object")
        return ret

169
    def _connect(self, other, connection, in_3prime_direction=None):
170
171
        ## TODO fix circular references        
        A,B = [connection.A, connection.B]
172
173
174
175
        if in_3prime_direction is not None:
            A.is_3prime_side_of_connection = not in_3prime_direction
            B.is_3prime_side_of_connection = in_3prime_direction
            
176
        A.connection = B.connection = connection
177
178
        self.connections.append(connection)
        other.connections.append(connection)
179
180
181
182
183
184
        l = A.container.locations
        if A not in l: l.append(A)
        l = B.container.locations
        if B not in l: l.append(B)
        

185
186
    # def _find_connections(self, loc):
    #     return [c for c in self.connections if c.A == loc or c.B == loc]
187
188
189

class SegmentParticle(PointParticle):
    def __init__(self, type_, position, name="A", segname="A", **kwargs):
190
        self.name = name
191
192
193
194
        self.contour_position = None
        PointParticle.__init__(self, type_, position, name=name, segname=segname, **kwargs)
        self.intrahelical_neighbors = []
        self.other_neighbors = []
cmaffeo2's avatar
cmaffeo2 committed
195
        self.locations = []
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210

    def get_intrahelical_above(self):
        """ Returns bead directly above self """
        assert( len(self.intrahelical_neighbors) <= 2 )
        for b in self.intrahelical_neighbors:
            if b.get_contour_position(self.parent) > self.contour_position:
                return b

    def get_intrahelical_below(self):
        """ Returns bead directly below self """
        assert( len(self.intrahelical_neighbors) <= 2 )
        for b in self.intrahelical_neighbors:
            if b.get_contour_position(self.parent) < self.contour_position:
                return b

cmaffeo2's avatar
cmaffeo2 committed
211
212
213
214
215
    def _neighbor_should_be_added(self,b):
        c1 = self.contour_position
        c2 = b.get_contour_position(self.parent)
        if c2 < c1:
            b0 = self.get_intrahelical_below()
216
        else:
cmaffeo2's avatar
cmaffeo2 committed
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
            b0 = self.get_intrahelical_above()

        if b0 is not None:            
            c0 = b0.get_contour_position(self.parent)
            if np.abs(c2-c1) < np.abs(c0-c1):
                ## remove b0
                self.intrahelical_neighbors.remove(b0)
                b0.intrahelical_neighbors.remove(self)
                return True
            else:
                return False
        return True
        
    def make_intrahelical_neighbor(self,b):
        add1 = self._neighbor_should_be_added(b)
        add2 = b._neighbor_should_be_added(self)
        if add1 and add2:
            assert(len(b.intrahelical_neighbors) <= 1)
            assert(len(self.intrahelical_neighbors) <= 1)
            self.intrahelical_neighbors.append(b)
            b.intrahelical_neighbors.append(self)
238

cmaffeo2's avatar
cmaffeo2 committed
239
240
241
242
243
244
        
    # def get_nt_position(self,seg):
    #     if seg == self.parent:
    #         return seg.contour_to_nt_pos(self.contour_position)
    #     else:
    #         cl = [e for e in self.parent.get_connections_and_locations() if e[2].container is seg]
245

cmaffeo2's avatar
cmaffeo2 committed
246
247
248
249
250
    #         dc = [(self.contour_position - A.address)**2 for c,A,B in cl]

    #         if len(dc) == 0:
    #             import pdb
    #             pdb.set_trace()
251

cmaffeo2's avatar
cmaffeo2 committed
252
253
254
255
256
257
258
259
260
261
    #         i = np.argmin(dc)
    #         c,A,B = cl[i]
    #         ## TODO: generalize, removing np.abs and conditional 
    #         delta_nt = np.abs( A.container.contour_to_nt_pos(self.contour_position - A.address) )
    #         B_nt_pos = seg.contour_to_nt_pos(B.address)
    #         if B.address < 0.5:
    #             return B_nt_pos-delta_nt
    #         else:
    #             return B_nt_pos+delta_nt

cmaffeo2's avatar
cmaffeo2 committed
262
    def get_nt_position_old(self, seg, nearest_position_on_seg=None):
263
        if seg == self.parent:
cmaffeo2's avatar
cmaffeo2 committed
264
            return seg.contour_to_nt_pos(self.contour_position)
265
266
        else:

cmaffeo2's avatar
cmaffeo2 committed
267
            def get_nt_pos(contour1, seg1, seg2):
cmaffeo2's avatar
cmaffeo2 committed
268
                cl = [e for e in seg1.get_connections_and_locations("intrahelical") if e[2].container is seg2]
cmaffeo2's avatar
cmaffeo2 committed
269
                dc = [(contour1 - A.address)**2 for c,A,B in cl]
cmaffeo2's avatar
cmaffeo2 committed
270
                if len(dc) == 0: return None,None
271

cmaffeo2's avatar
cmaffeo2 committed
272
273
274
275
                i = np.argmin(dc)
                c,A,B = cl[i]

                ## TODO: generalize, removing np.abs and conditional 
cmaffeo2's avatar
cmaffeo2 committed
276
277
                # delta_nt = np.abs( seg1.contour_to_nt_pos(contour1 - A.address) )
                delta_nt = seg1.contour_to_nt_pos(contour1) - seg1.contour_to_nt_pos(A.address)
cmaffeo2's avatar
cmaffeo2 committed
278
                B_nt_pos = seg2.contour_to_nt_pos(B.address)
cmaffeo2's avatar
cmaffeo2 committed
279
280
281
282
283
284
285
                dirA = 2*(A.address > 0.5)-1
                dirB = 2*(B.address < 0.5)-1
                direction = dirA*dirB
                if c.type_ in ("intrahelical",):
                    pass
                elif c.type_ in ("crossover","terminal-crossover"):
                    direction *= -1
cmaffeo2's avatar
cmaffeo2 committed
286
                else:
cmaffeo2's avatar
cmaffeo2 committed
287
288
289
                    raise Exception("Unhandled connection type {}".format(c.type_))
                distance = np.abs(delta_nt)+np.abs(B_nt_pos)
                return B_nt_pos + direction*delta_nt, distance
cmaffeo2's avatar
cmaffeo2 committed
290
                
cmaffeo2's avatar
cmaffeo2 committed
291
            pos,dist = get_nt_pos(self.contour_position, self.parent, seg)
cmaffeo2's avatar
cmaffeo2 committed
292
293
294
295
296
297
            if pos is None:
                ## Particle is not directly connected
                visited_segs = set(seg)
                positions = []
                for l in self.locations:
                    if l.container == self.parent: continue
cmaffeo2's avatar
cmaffeo2 committed
298
299
                    if l.connection.type_ != "intrahelical": continue
                    pos0,dist = get_nt_pos(self.contour_position, self.parent, l.container)
cmaffeo2's avatar
cmaffeo2 committed
300
301
                    assert(pos0 is not None)
                    pos0 = l.container.nt_pos_to_contour(pos0)
cmaffeo2's avatar
cmaffeo2 committed
302
                    pos,dist = get_nt_pos( pos0, l.container, seg )
cmaffeo2's avatar
cmaffeo2 committed
303
304
305
306
307
308
309
310
                    if pos is not None:
                        positions.append( pos )
                assert( len(positions) > 0 )
                if len(positions) > 1:
                    import pdb
                    pdb.set_trace()
                pos = positions[0]
            return pos
311

cmaffeo2's avatar
cmaffeo2 committed
312
313
314
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
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
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
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
    def get_nt_position(self, seg):
        """ Returns the "address" of the nucleotide relative to seg in
        nucleotides, taking the shortest (intrahelical) contour length route to seg
        """

        if seg == self.parent:
            return seg.contour_to_nt_pos(self.contour_position)
        else:
            pos = self.get_contour_position(seg)
            return seg.contour_to_nt_pos(pos)


            # def get_nt_pos(contour1, seg1, seg2):
            #     cl = [e for e in seg1.get_connections_and_locations("intrahelical") if e[2].container is seg2]
            #     dc = [(contour1 - A.address)**2 for c,A,B in cl]
            #     if len(dc) == 0: return None,None

            #     i = np.argmin(dc)
            #     c,A,B = cl[i]

            #     ## TODO: generalize, removing np.abs and conditional 
            #     # delta_nt = np.abs( seg1.contour_to_nt_pos(contour1 - A.address) )
            #     delta_nt = seg1.contour_to_nt_pos(contour1) - seg1.contour_to_nt_pos(A.address)
            #     B_nt_pos = seg2.contour_to_nt_pos(B.address)
            #     dirA = 2*(A.address > 0.5)-1
            #     dirB = 2*(B.address < 0.5)-1
            #     direction = dirA*dirB
            #     distance = np.abs(delta_nt)+np.abs(B_nt_pos)
            #     return B_nt_pos + direction*delta_nt, distance
            # pos,dist = get_nt_pos(self.contour_position, self.parent, seg)

            cutoff = 30*3
            target_seg = seg

            ## depth-first search
            ## TODO cache distances to nearby locations?
            def descend_search_tree(seg, contour_in_seg, distance=0, visited_segs=[]):
                nonlocal cutoff

                if seg == target_seg:
                    # pdb.set_trace()
                    ## Found a segment in our target
                    sign = (contour_in_seg == 1) - (contour_in_seg == 0)
                    assert( sign in (1,-1) )
                    if distance < cutoff: # TODO: check if this does anything
                        cutoff = distance
                    return [[distance, contour_in_seg+sign*seg.nt_pos_to_contour(distance)]]
                if distance > cutoff:
                    return None
                    
                ret_list = []
                ## Find intrahelical locations in seg that we might pass through
                for c,A,B in seg.get_connections_and_locations("intrahelical"):
                    if B.container in visited_segs: continue
                    dx = seg.contour_to_nt_pos( np.abs(A.address-contour_in_seg),
                                                round_nt=False) 
                    results = descend_search_tree( B.container, B.address,
                                                   distance+dx, visited_segs + [seg] )
                    if results is not None:
                        ret_list.extend( results )
                return ret_list

            results = descend_search_tree(self.parent, self.contour_position)
            if results is None or len(results) == 0:
                raise Exception("Could not find location in segment") # TODO better error
            return sorted(results,key=lambda x:x[0])[0][1]

            # def descend_search_tree(seg, target_seg, distance, from_contour):
            #     if distance > cutoff:
            #         return None
            #     ret_list = []

            #     ## Find intrahelical locations in seg that we might pass through
            #     for c,A,B in seg.get_connections_and_locations("intrahelical"):
            #         if np.isclose(A.address, from_contour): continue
            #         if B.container == target_seg:
            #             ## Found a segment in our target
            #             ...
            #         else:
            #             ret_list.append( [B.container, seg.num_nts] ]
                

            # curr_seg = seg
            # dist = 0
            # for c,A,B in curr_seg.get_connections_and_locations("intrahelical"):
            #     while 

            # cl = [e for e in seg1.get_connections_and_locations("intrahelical") if e[2].container is seg2]



            # if dist is not None and dist < closest_distance:
            #     closest_distance
            # if pos is None:
            #     ## Particle is not directly connected
            #     visited_segs = set(seg)
            #     positions = []
            #     for l in self.locations:
            #         if l.container == self.parent: continue
            #         if l.connection.type_ != "intrahelical": continue
            #         pos0,dist = get_nt_pos(self.contour_position, self.parent, l.container)
            #         assert(pos0 is not None)
            #         pos0 = l.container.nt_pos_to_contour(pos0)
            #         pos,dist = get_nt_pos( pos0, l.container, seg )
            #         if pos is not None:
            #             positions.append( pos )
            #     assert( len(positions) > 0 )
            #     if len(positions) > 1:
            #         import pdb
            #         pdb.set_trace()
            #     pos = positions[0]
            # return pos

425

426
427
428
429
    def get_contour_position(self,seg):
        if seg == self.parent:
            return self.contour_position
        else:
cmaffeo2's avatar
cmaffeo2 committed
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
            cutoff = 30*3
            target_seg = seg

            ## depth-first search
            ## TODO cache distances to nearby locations?
            def descend_search_tree(seg, contour_in_seg, distance=0, visited_segs=[]):
                nonlocal cutoff

                if seg == target_seg:
                    # pdb.set_trace()
                    ## Found a segment in our target
                    sign = (contour_in_seg == 1) - (contour_in_seg == 0)
                    assert( sign in (1,-1) )
                    if distance < cutoff: # TODO: check if this does anything
                        cutoff = distance
                    return [[distance, contour_in_seg+sign*seg.nt_pos_to_contour(distance)]]
                if distance > cutoff:
                    return None
                    
                ret_list = []
                ## Find intrahelical locations in seg that we might pass through
                for c,A,B in seg.get_connections_and_locations("intrahelical"):
                    if B.container in visited_segs: continue
                    dx = seg.contour_to_nt_pos( np.abs(A.address-contour_in_seg),
                                                round_nt=False) 
                    results = descend_search_tree( B.container, B.address,
                                                   distance+dx, visited_segs + [seg] )
                    if results is not None:
                        ret_list.extend( results )
                return ret_list

            results = descend_search_tree(self.parent, self.contour_position)
            if results is None or len(results) == 0:
                raise Exception("Could not find location in segment") # TODO better error
            return sorted(results,key=lambda x:x[0])[0][1]

            # nt_pos = self.get_nt_position(seg)
            # return seg.nt_pos_to_contour(nt_pos)

    def update_position(self, contour_position):
        self.contour_position = contour_position
        self.position = self.parent.contour_to_position(contour_position)
        if 'orientation_bead' in self.__dict__:
            o = self.orientation_bead
            o.contour_position = contour_position
            orientation = self.parent.contour_to_orientation(contour_position)
            if orientation is None:
                print("WARNING: local_twist is True, but orientation is None; using identity")
                orientation = np.eye(3)
            o.position = self.position + orientation.dot( np.array((Segment.orientation_bond.r0,0,0)) )
            
481
482

## TODO break this class into smaller, better encapsulated pieces
483
484
485
486
487
488
489
490
491
492
493
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
494
495
496
497
498
    orientation_particle = ParticleType("O",
                                        diffusivity = 100,
                                        mass = 300,
                                        radius = 1,
                                    )
499

cmaffeo2's avatar
cmaffeo2 committed
500
    # orientation_bond = HarmonicBond(10,2)
501
    orientation_bond = HarmonicBond(30,1.5, rRange = (0,500) )
502
503
504
505
506
507
508
509
510
511
512
513
514

    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=[])
515
        ConnectableElement.__init__(self, connection_locations=[], connections=[])
516

517
        self.resname = name
cmaffeo2's avatar
cmaffeo2 committed
518
519
520
521
522
        self.start_orientation = None
        self.twist_per_nt = 0

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

523
524
525
        self._bead_model_generation = 0    # TODO: remove?
        self.segment_model = segment_model # TODO: remove?

cmaffeo2's avatar
cmaffeo2 committed
526
        self.num_nts = int(num_nts)
527
528
529
530
531
        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

532
533
534
535
        ## 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
536
537
        
        self.sequence = None
538

539
540
541
542
543
    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
544

545
    def contour_to_nt_pos(self, contour_pos, round_nt=False):
cmaffeo2's avatar
cmaffeo2 committed
546
        nt = contour_pos*(self.num_nts) - 0.5
547
        if round_nt:
cmaffeo2's avatar
cmaffeo2 committed
548
            assert( np.isclose(np.around(nt),nt) )
549
550
551
            nt = np.around(nt)
        return nt

552
    def nt_pos_to_contour(self,nt_pos):
cmaffeo2's avatar
cmaffeo2 committed
553
        return (nt_pos+0.5)/(self.num_nts)
554

555
556
557
558
559
560
561
    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 )
562
563
        t = (t / np.linalg.norm(t,axis=0))
        return t.T
564
565
566
        

    def contour_to_orientation(self,s):
567
568
        assert( isinstance(s,float) or isinstance(s,int) or len(s) == 1 )   # TODO make vectorized version
        orientation = None
569
570
571
572
        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)
573
                orientation = rotationAboutAxis( axis, self.twist_per_nt*self.contour_to_nt_pos(s), normalizeAxis=True )
574
575
576
577
578
            else:
                q = interpolate.splev( s, self.quaternion_spline_params )
                if len(q) > 1: q = np.array(q).T # TODO: is this needed?
                orientation = quaternion_to_matrix(q)
        return orientation
579

cmaffeo2's avatar
cmaffeo2 committed
580
    def get_contour_sorted_connections_and_locations(self,type_):
cmaffeo2's avatar
cmaffeo2 committed
581
        sort_fn = lambda c: c[1].address
cmaffeo2's avatar
cmaffeo2 committed
582
        cl = self.get_connections_and_locations(type_)
cmaffeo2's avatar
cmaffeo2 committed
583
        return sorted(cl, key=sort_fn)
584
585
586
    
    def randomize_unset_sequence(self):
        bases = list(seqComplement.keys())
587
        # bases = ['T']        ## FOR DEBUG
588
589
590
591
592
593
594
        if self.sequence is None:
            self.sequence = [random.choice(bases) for i in range(self.num_nts)]
        else:
            assert(len(self.sequence) == self.num_nts) # TODO move
            for i in range(len(self.sequence)):
                if self.sequence[i] is None:
                    self.sequence[i] = random.choice(bases)
595

cmaffeo2's avatar
cmaffeo2 committed
596
597
598
    def _get_num_beads(self, max_basepairs_per_bead, max_nucleotides_per_bead ):
        raise NotImplementedError

599
    def _generate_one_bead(self, contour_position, nts):
600
601
        raise NotImplementedError

602
    def _generate_atomic_nucleotide(self, contour_position, is_fwd, seq, scale):
cmaffeo2's avatar
cmaffeo2 committed
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
        """ 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
619
                ## TODO: improve placement of ssDNA
cmaffeo2's avatar
cmaffeo2 committed
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
                # 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
        if not is_fwd:
            nt_dict = canonicalNtFwd
        else:
            nt_dict = canonicalNtRev
638
        atoms = nt_dict[ key ].generate() # TODO: clone?
cmaffeo2's avatar
cmaffeo2 committed
639
                        
cmaffeo2's avatar
cmaffeo2 committed
640
        atoms.orientation = orientation.dot(atoms.orientation)
641
642
643
644
645
        if isinstance(self, SingleStrandedSegment):
            if scale is not None and scale != 1:
                for a in atoms:
                    a.position = scale*a.position
                    a.beta = 0
646
            atoms.position = pos - atoms.atoms_by_name["C1'"].collapsedPosition()
647
648
649
650
651
652
653
654
655
656
657
        else:
            if scale is not None and scale != 1:
                if atoms.sequence in ("A","G"):
                    r0 = atoms.atoms_by_name["N9"].position
                else:
                    r0 = atoms.atoms_by_name["N1"].position
                for a in atoms:
                    if a.name[-1] in ("'","P","T"):
                        a.position = scale*(a.position-r0) + r0
                        a.beta = 0
            atoms.position = pos
cmaffeo2's avatar
cmaffeo2 committed
658
659

        return atoms
660

661
662
    def add_location(self, nt, type_, on_fwd_strand=True):
        ## Create location if needed, add to segment
663
        c = self.nt_pos_to_contour(nt)
664
665
666
667
668
669
670
        assert(c >= 0 and c <= 1)
        # TODO? loc = self.Location( address=c, type_=type_, on_fwd_strand=is_fwd )
        loc = Location( self, address=c, type_=type_, on_fwd_strand=on_fwd_strand )
        self.locations.append(loc)

    ## TODO? Replace with abstract strand-based model?
    def add_5prime(self, nt, on_fwd_strand=True):
671
672
        if isinstance(self,SingleStrandedSegment):
            on_fwd_strand = True
673
        self.add_location(nt,"5prime",on_fwd_strand)
674
675

    def add_3prime(self, nt, on_fwd_strand=True):
676
677
        if isinstance(self,SingleStrandedSegment):
            on_fwd_strand = True
678
        self.add_location(nt,"3prime",on_fwd_strand)
679

680
    def get_3prime_locations(self):
cmaffeo2's avatar
cmaffeo2 committed
681
        return sorted(self.get_locations("3prime"),key=lambda x: x.address)
682
    
cmaffeo2's avatar
cmaffeo2 committed
683
    def get_5prime_locations(self):
684
        ## TODO? ensure that data is consistent before _build_model calls
cmaffeo2's avatar
cmaffeo2 committed
685
        return sorted(self.get_locations("5prime"),key=lambda x: x.address)
cmaffeo2's avatar
cmaffeo2 committed
686

687
    def iterate_connections_and_locations(self, reverse=False):
cmaffeo2's avatar
cmaffeo2 committed
688
689
        ## connections to other segments
        cl = self.get_contour_sorted_connections_and_locations()
690
        if reverse:
cmaffeo2's avatar
cmaffeo2 committed
691
            cl = cl[::-1]
692
693
694
            
        for c in cl:
            yield c
cmaffeo2's avatar
cmaffeo2 committed
695

696
    ## TODO rename
697
    def get_strand_segment(self, nt_pos, is_fwd, move_at_least=0.5):
698
        """ Walks through locations, checking for crossovers """
699
700
701
702
        # if self.name in ("6-1","1-1"):
        #     import pdb
        #     pdb.set_trace()
        move_at_least = 0
703
704

        ## Iterate through locations
cmaffeo2's avatar
cmaffeo2 committed
705
        # locations = sorted(self.locations, key=lambda l:(l.address,not l.on_fwd_strand), reverse=(not is_fwd))
706
707
708
709
710
711
        def loc_rank(l):
            nt = l.get_nt_pos()
            ## optionally add logic about type of connection
            return (nt, not l.on_fwd_strand)
        # locations = sorted(self.locations, key=lambda l:(l.address,not l.on_fwd_strand), reverse=(not is_fwd))
        locations = sorted(self.locations, key=loc_rank, reverse=(not is_fwd))
712
713
        # print(locations)

714
        for l in locations:
cmaffeo2's avatar
cmaffeo2 committed
715
716
717
718
719
720
721
            # TODOTODO probably okay
            if l.address == 0:
                pos = 0.0
            elif l.address == 1:
                pos = self.num_nts-1
            else:
                pos = self.contour_to_nt_pos(l.address, round_nt=True)
722
723
724

            ## DEBUG

cmaffeo2's avatar
cmaffeo2 committed
725

726
            ## Skip locations encountered before our strand
727
728
729
730
731
732
733
734
            # tol = 0.1
            # if is_fwd:
            #     if pos-nt_pos <= tol: continue 
            # elif   nt_pos-pos <= tol: continue
            if (pos-nt_pos)*(2*is_fwd-1) < move_at_least: continue
            ## TODO: remove move_at_least
            if np.isclose(pos,nt_pos):
                if l.is_3prime_side_of_connection: continue
735
736
737

            ## Stop if we found the 3prime end
            if l.on_fwd_strand == is_fwd and l.type_ == "3prime":
738
739
                print("  found end at",l)
                return pos, None, None, None, None
740
741
742
743
744
745
746
747

            ## Check location connections
            c = l.connection
            if c is None: continue
            B = c.other(l)            

            ## Found a location on the same strand?
            if l.on_fwd_strand == is_fwd:
748
749
                print("  passing through",l)
                print("from {}, connection {} to {}".format(nt_pos,l,B))
750
                Bpos = B.get_nt_pos()
751
                return pos, B.container, Bpos, B.on_fwd_strand, 0.5
752
753
754
                
            ## Stop at other strand crossovers so basepairs line up
            elif c.type_ == "crossover":
755
756
757
                if nt_pos == pos: continue
                print("  pausing at",l)
                return pos, l.container, pos+(2*is_fwd-1), is_fwd, 0
758

759
760
        import pdb
        pdb.set_trace()
761
762
763
764
765
        raise Exception("Shouldn't be here")
        # print("Shouldn't be here")
        ## Made it to the end of the segment without finding a connection
        return 1*is_fwd, None, None, None

766
767
768
    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
769
        # TODO: include beads in connections?
770
771
772
        i = np.argmin((cs - contour_position)**2)

        return self.beads[i]
773
774
775

    def get_all_consecutive_beads(self, number):
        assert(number >= 1)
cmaffeo2's avatar
cmaffeo2 committed
776
        ## Assume that consecutive beads in self.beads are bonded
777
        ret = []
cmaffeo2's avatar
cmaffeo2 committed
778
779
        for i in range(len(self.beads)-number+1):
            tmp = [self.beads[i+j] for j in range(0,number)]
780
            ret.append( tmp )
781
        return ret   
782

783
784
785
    def _add_bead(self,b,set_contour=False):
        if set_contour:
            b.contour_position = b.get_contour_position(self)
786
        
787
788
789
        # assert(b.parent is None)
        if b.parent is not None:
            b.parent.children.remove(b)
790
        self.add(b)
791
792
793
794
795
796
        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)
797
            self.add(o)
798
799
800
801
802
803
804
805
            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 = []
806
807
808

        if True:
            print("WARNING: DEBUG")
809
            ## Remove duplicates, preserving order
810
811
812
813
            tmp = []
            for c in new_children:
                if c not in tmp:
                    tmp.append(c)
814
815
                else:
                    print("  duplicate particle found!")
816
817
            new_children = tmp

818
819
820
821
822
        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)
823
824
825
826
827
            
        # tmp = [c for c in self.children if c not in old_children]
        # assert(len(tmp) == 0)
        # tmp = [c for c in old_children if c not in self.children]
        # assert(len(tmp) == 0)
828
829
        assert(len(old_children) == len(self.children))
        assert(len(old_beads) == len(self.beads))
830

831

cmaffeo2's avatar
cmaffeo2 committed
832
    def _generate_beads(self, bead_model, max_basepairs_per_bead, max_nucleotides_per_bead):
833

834
        """ Generate beads (positions, types, etc) and bonds, angles, dihedrals, exclusions """
cmaffeo2's avatar
cmaffeo2 committed
835
        ## TODO: decide whether to remove bead_model argument
836
        ##       (currently unused)
cmaffeo2's avatar
cmaffeo2 committed
837

838
        ## First find points between-which beads must be generated
839
840
841
        # conn_locs = self.get_contour_sorted_connections_and_locations()
        # locs = [A for c,A,B in conn_locs]
        # existing_beads = [l.particle for l in locs if l.particle is not None]
cmaffeo2's avatar
cmaffeo2 committed
842
843
        # if self.name == "31-2":
            # pdb.set_trace()
844

845
846
847
848
849
        existing_beads = {l.particle for l in self.locations if l.particle is not None}
        existing_beads = sorted( list(existing_beads), key=lambda b: b.get_contour_position(self) )
        
        if len(existing_beads) != len(set(existing_beads)):
            pdb.set_trace()
850
851
852
        for b in existing_beads:
            assert(b.parent is not None)

cmaffeo2's avatar
cmaffeo2 committed
853
        # if self.name in ("31-2",):
cmaffeo2's avatar
cmaffeo2 committed
854
855
        #     pdb.set_trace()

856
        ## Add ends if they don't exist yet
857
        ## TODOTODO: test 1 nt segments?
cmaffeo2's avatar
cmaffeo2 committed
858
859
860
        if len(existing_beads) == 0 or existing_beads[0].get_nt_position(self) > 0.5:
            # if len(existing_beads) > 0:            
            #     assert(existing_beads[0].get_nt_position(self) >= 0.5)
861
862
            b = self._generate_one_bead(0, 0)
            existing_beads = [b] + existing_beads
cmaffeo2's avatar
cmaffeo2 committed
863
864

        if existing_beads[-1].get_nt_position(self)-(self.num_nts-1) < -0.5:
865
866
867
868
869
870
            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
871
        last = None
cmaffeo2's avatar
cmaffeo2 committed
872

873
874
        for I in range(len(existing_beads)-1):
            eb1,eb2 = [existing_beads[i] for i in (I,I+1)]
875
            assert( eb1 is not eb2 )
876

cmaffeo2's avatar
cmaffeo2 committed
877
878
879
880
881
            # if np.isclose(eb1.position[2], eb2.position[2]):
            #     import pdb
            #     pdb.set_trace()

            print(" %s working on %d to %d" % (self.name, eb1.position[2], eb2.position[2]))
882
883
884
885
886
887
888
889
890
891
892
893
894
            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:
cmaffeo2's avatar
cmaffeo2 committed
895
                last.make_intrahelical_neighbor(eb1)
896
897
898
            last = eb1
            for j in range(num_beads):
                s = ds*(j+1) + s0
899
                # if self.name in ("51-2","51-3"):
cmaffeo2's avatar
cmaffeo2 committed
900
                # if self.name in ("31-2",):
901
                #     print(" adding bead at {}".format(s))
902
903
                b = self._generate_one_bead(s,nts)

cmaffeo2's avatar
cmaffeo2 committed
904
                last.make_intrahelical_neighbor(b)
905
906
907
                last = b
                tmp_children.append(b)

cmaffeo2's avatar
cmaffeo2 committed
908
        last.make_intrahelical_neighbor(eb2)
909
910
911

        if eb2.parent == self:
            tmp_children.append(eb2)
cmaffeo2's avatar
cmaffeo2 committed
912
        # if self.name in ("31-2",):
913
        #     pdb.set_trace()
914
        self._rebuild_children(tmp_children)
915
916
917
918
919
920
921
922
923
924
925
926
927

    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
928
929
                 local_twist = False,
                 num_turns = None,
cmaffeo2's avatar
cmaffeo2 committed
930
931
                 start_orientation = None,
                 twist_persistence_length = 90 ):
cmaffeo2's avatar
cmaffeo2 committed
932
933
934
        
        self.helical_rise = 10.44
        self.distance_per_nt = 3.4
935
936
937
938
939
        Segment.__init__(self, name, num_nts, 
                         start_position,
                         end_position, 
                         segment_model)

cmaffeo2's avatar
cmaffeo2 committed
940
941
942
943
944
945
        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:
946
            start_orientation = np.eye(3) # np.array(((1,0,0),(0,1,0),(0,0,1)))
cmaffeo2's avatar
cmaffeo2 committed
947
        self.start_orientation = start_orientation
cmaffeo2's avatar
cmaffeo2 committed
948
        self.twist_persistence_length = twist_persistence_length
cmaffeo2's avatar
cmaffeo2 committed
949

950
951
        self.nicks = []

952
        self.start = self.start5 = Location( self, address=0, type_= "end5" )
953
        self.start3 = Location( self, address=0, type_ = "end3", on_fwd_strand=False )
954

955
956
        self.end = self.end3 = Location( self, address=1, type_ = "end3" )
        self.end5 = Location( self, address=1, type_= "end5", on_fwd_strand=False )
cmaffeo2's avatar
cmaffeo2 committed
957
958
        # for l in (self.start5,self.start3,self.end3,self.end5):
        #     self.locations.append(l)
959

960
961
962
963
964
965
966
967
968
        ## 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


969
    ## Convenience methods
970
    ## TODO: add errors if unrealistic connections are made
971
    ## TODO: make connections automatically between unconnected strands
972
    def connect_start5(self, end3, type_="intrahelical", force_connection=False):
973
974
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
975
976
        self._connect_ends( self.start5, end3, type_, force_connection = force_connection )
    def connect_start3(self, end5, type_="intrahelical", force_connection=False):
977
        if isinstance(end5, SingleStrandedSegment):
978
            end5 = end5.start5
979
980
        self._connect_ends( self.start3, end5, type_, force_connection = force_connection )
    def connect_end3(self, end5, type_="intrahelical", force_connection=False):
981
        if isinstance(end5, SingleStrandedSegment):
982
            end5 = end5.start5
983
984
        self._connect_ends( self.end3, end5, type_, force_connection = force_connection )
    def connect_end5(self, end3, type_="intrahelical", force_connection=False):
985
986
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
987
        self._connect_ends( self.end5, end3, type_, force_connection = force_connection )
988

cmaffeo2's avatar
cmaffeo2 committed
989
    def add_crossover(self, nt, other, other_nt, strands_fwd=[True,False], nt_on_5prime=True):
cmaffeo2's avatar
cmaffeo2 committed
990
991
992
993
        """ Add a crossover between two helices """
        ## Validate other, nt, other_nt
        ##   TODO

994
        if isinstance(other,SingleStrandedSegment):
cmaffeo2's avatar
cmaffeo2 committed
995
            other.add_crossover(other_nt, self, nt, strands_fwd[::-1], not nt_on_5prime)
996
        else:
997

998
999
1000
            ## Create locations, connections and add to segments
            c = self.nt_pos_to_contour(nt)
            assert(c >= 0 and c <= 1)
For faster browsing, not all history is shown. View entire blame