segmentmodel.py 102 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
                # print("  found end at",l)
739
                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
                if nt_pos == pos: continue
756
                # print("  pausing at",l)
757
                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

        if True:
808
809
            ## TODO: remove this if duplicates are never found 
            # print("Searching for duplicate particles...")
810
            ## Remove duplicates, preserving order
811
812
813
814
            tmp = []
            for c in new_children:
                if c not in tmp:
                    tmp.append(c)
815
                else:
816
                    print("  DUPLICATE PARTICLE FOUND!")
817
818
            new_children = tmp

819
820
821
822
823
        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)
824
825
826
827
828
            
        # 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)
829
830
        assert(len(old_children) == len(self.children))
        assert(len(old_beads) == len(self.beads))
831

832

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

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

839
        ## First find points between-which beads must be generated
840
841
842
        # 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
843
844
        # if self.name == "31-2":
            # pdb.set_trace()
845

846
847
848
849
850
        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()
851
852
853
        for b in existing_beads:
            assert(b.parent is not None)

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

857
        ## Add ends if they don't exist yet
858
        ## TODOTODO: test 1 nt segments?
cmaffeo2's avatar
cmaffeo2 committed
859
860
861
        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)
862
            b = self._generate_one_bead( self.nt_pos_to_contour(0), 0)
863
            existing_beads = [b] + existing_beads
cmaffeo2's avatar
cmaffeo2 committed
864
865

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

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

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

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

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

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

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

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

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

951
952
        self.nicks = []

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

956
957
        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
958
959
        # for l in (self.start5,self.start3,self.end3,self.end5):
        #     self.locations.append(l)
960

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


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

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

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

999
1000
1001
            ## Create locations, connections and add to segments
            c = self.nt_pos_to_contour(nt)
            assert(c >= 0 and c <= 1)
1002

1003
1004
1005
            loc = self.get_location_at(c, strands_fwd[0])

            c = other.nt_pos_to_contour(other_nt)
cmaffeo2's avatar
cmaffeo2 committed
1006
            # TODOTODO: may need to subtract or add a little depending on 3prime/5prime
1007
1008
1009
            assert(c >= 0 and c <= 1)
            other_loc = other.get_location_at(c, strands_fwd[1])
            self._connect(other, Connection( loc, other_loc, type_="crossover" ))
cmaffeo2's avatar
cmaffeo2 committed
1010
1011
1012
1013
1014
1015
            if nt_on_5prime:
                loc.is_3prime_side_of_connection = False
                other_loc.is_3prime_side_of_connection = True
            else:            
                loc.is_3prime_side_of_connection = True
                other_loc.is_3prime_side_of_connection = False
cmaffeo2's avatar
cmaffeo2 committed
1016

1017
    ## Real work
1018
    def _connect_ends(self, end1, end2, type_, force_connection):
1019
        ## TODO remove self?
1020
1021
1022
1023
1024
        ## validate the input
        for end in (end1, end2):
            assert( isinstance(end, Location) )
            assert( end.type_ in ("end3","end5") )
        assert( end1.type_ != end2.type_ )
1025
        ## Create and add connection
cmaffeo2's avatar
cmaffeo2 committed
1026
        if end2.type_ == "end5":
1027
1028
1029
            end1.container._connect( end2.container, Connection( end1, end2, type_=type_ ), in_3prime_direction=True )
        else:
            end2.container._connect( end1.container, Connection( end2, end1, type_=type_ ), in_3prime_direction=True )
1030
    def _get_num_beads(self, contour, max_basepairs_per_bead, max_nucleotides_per_bead):
1031
1032
        # return int(contour*self.num_nts // max_basepairs_per_bead)
        return int(contour*(self.num_nts**2/(self.num_nts+1)) // max_basepairs_per_bead)
cmaffeo2's avatar
cmaffeo2 committed
1033

1034
1035
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
1036
        if self.local_twist:
1037
            orientation = self.contour_to_orientation(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
1038
1039
1040
            if orientation is None:
                print("WARNING: local_twist is True, but orientation is None; using identity")
                orientation = np.eye(3)
cmaffeo2's avatar
cmaffeo2 committed
1041
            opos = pos + orientation.dot( np.array((Segment.orientation_bond.r0,0,0)) )
cmaffeo2's avatar
cmaffeo2 committed
1042
1043
1044
1045
1046
            # if np.linalg.norm(pos) > 1e3:
            #     pdb.set_trace()
            assert(np.linalg.norm(opos-pos) < 10 )
            o = SegmentParticle( Segment.orientation_particle, opos, name="O",
                                 contour_position = contour_position,
1047
                                 num_nts=nts, parent=self )
1048
            bead = SegmentParticle( Segment.dsDNA_particle, pos, name="DNA",
1049
1050
1051
                                    num_nts=nts, parent=self, 
                                    orientation_bead=o,
                                    contour_position=contour_position )
cmaffeo2's avatar
cmaffeo2 committed
1052
1053

        else:
1054
            bead = SegmentParticle( Segment.dsDNA_particle, pos, name="DNA",
1055
1056
1057
                                    num_nts=nts, parent=self,
                                    contour_position=contour_position )
        self._add_bead(bead)
cmaffeo2's avatar
cmaffeo2 committed
1058
        return bead
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074

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)

1075
        self.start = self.start5 = Location( self, address=0, type_= "end5" ) # TODO change type_?
1076
        self.end = self.end3 = Location( self, address=1, type_ = "end3" )
cmaffeo2's avatar
cmaffeo2 committed
1077
1078
        # for l in (self.start5,self.end3):
        #     self.locations.append(l)
1079

1080
    def connect_end3(self, end5, force_connection=False):
cmaffeo2's avatar
cmaffeo2 committed
1081
        self._connect_end( end5,  _5_to_3 = True, force_connection = force_connection )
1082

1083
    def connect_5end(self, end3, force_connection=False): # TODO: change name or possibly deprecate
cmaffeo2's avatar
cmaffeo2 committed
1084
        self._connect_end( end3,  _5_to_3 = False, force_connection = force_connection )
1085
1086
1087
1088

    def _connect_end(self, other, _5_to_3, force_connection):
        assert( isinstance(other, Location) )
        if _5_to_3 == True:
cmaffeo2's avatar
cmaffeo2 committed
1089
1090
1091
1092
            my_end = self.end3
            # assert( other.type_ == "end5" )
            if (other.type_ is not "end5"):
                print("Warning: code does not prevent connecting 3prime to 3prime, etc")
1093
1094
            conn = Connection( my_end, other, type_="intrahelical" )
            self._connect( other.container, conn, in_3prime_direction=True )
1095
        else:
cmaffeo2's avatar
cmaffeo2 committed
1096
1097
1098
1099
            my_end = self.end5
            # assert( other.type_ == "end3" )
            if (other.type_ is not "end3"):
                print("Warning: code does not prevent connecting 3prime to 3prime, etc")
1100
1101
            conn = Connection( other, my_end, type_="intrahelical" )
            other.container._connect( self, conn, in_3prime_direction=True )
1102

cmaffeo2's avatar
cmaffeo2 committed
1103
    def add_crossover(self, nt, other, other_nt, strands_fwd=[True,False], nt_on_5prime=True):
1104
1105
1106
1107
1108
1109
        """ Add a crossover between two helices """
        ## Validate other, nt, other_nt
        ##   TODO
       
        ## TODO: fix direction

cmaffeo2's avatar
cmaffeo2 committed
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
        # c1 = self.nt_pos_to_contour(nt)
        # # TODOTODO
        # ## Ensure connections occur at ends, otherwise the structure doesn't make sense
        # # assert(np.isclose(c1,0) or np.isclose(c1,1))
        # assert(np.isclose(nt,0) or np.isclose(nt,self.num_nts-1))
        if nt == 0:
            c1 = 0