segmentmodel.py 165 KB
Newer Older
1001
1002
                if _DEBUG_TRACE:
                    print("  found end at",l)
1003
                return pos, None, None, None, None
1004
1005
1006
1007
1008
1009
1010
1011

            ## 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:
1012
1013
1014
                if _DEBUG_TRACE:
                    print("  passing through",l)
                    print("from {}, connection {} to {}".format(nt_pos,l,B))
1015
                Bpos = B.get_nt_pos()
1016
                return pos, B.container, Bpos, B.on_fwd_strand, 0.5
1017
1018
1019
                
            ## Stop at other strand crossovers so basepairs line up
            elif c.type_ == "crossover":
1020
                if nt_pos == pos: continue
1021
1022
                if _DEBUG_TRACE:
                    print("  pausing at",l)
1023
                return pos, l.container, pos+(2*is_fwd-1), is_fwd, 0
1024
1025
1026
1027
1028
1029

        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

1030
1031
1032
    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
1033
        # TODO: include beads in connections?
1034
1035
1036
        i = np.argmin((cs - contour_position)**2)

        return self.beads[i]
1037

1038
1039
1040
1041
1042
1043
1044
1045
1046
    def _get_atomic_nucleotide(self, nucleotide_idx, is_fwd=True):
        d = 'fwd' if is_fwd else 'rev'
        for s in self.strand_pieces[d]:
            try:
                return s.get_nucleotide(nucleotide_idx)
            except:
                pass
        raise Exception("Could not find nucleotide in {} at {}.{}".format( self, nucleotide_idx, d ))

1047
1048
    def get_all_consecutive_beads(self, number):
        assert(number >= 1)
cmaffeo2's avatar
cmaffeo2 committed
1049
        ## Assume that consecutive beads in self.beads are bonded
1050
        ret = []
cmaffeo2's avatar
cmaffeo2 committed
1051
1052
        for i in range(len(self.beads)-number+1):
            tmp = [self.beads[i+j] for j in range(0,number)]
1053
            ret.append( tmp )
1054
        return ret   
1055

1056
    def _add_bead(self,b):
1057
        
1058
1059
1060
        # assert(b.parent is None)
        if b.parent is not None:
            b.parent.children.remove(b)
1061
        self.add(b)
1062
1063
1064
1065
1066
1067
        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)
1068
            self.add(o)
1069
1070
1071
1072
1073
1074
1075
1076
            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 = []
1077
1078

        if True:
1079
1080
            ## TODO: remove this if duplicates are never found 
            # print("Searching for duplicate particles...")
1081
            ## Remove duplicates, preserving order
1082
1083
1084
1085
            tmp = []
            for c in new_children:
                if c not in tmp:
                    tmp.append(c)
1086
                else:
1087
                    print("  DUPLICATE PARTICLE FOUND!")
1088
1089
            new_children = tmp

1090
1091
1092
1093
1094
        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)
1095
1096
1097
1098
1099
            
        # 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)
1100
1101
        assert(len(old_children) == len(self.children))
        assert(len(old_beads) == len(self.beads))
1102

1103

1104
    def _generate_beads(self, bead_model, max_basepairs_per_bead, max_nucleotides_per_bead, one_bead_per_monomer=False):
1105
        """ Generate beads (positions, types, etc) and bonds, angles, dihedrals, exclusions """
cmaffeo2's avatar
cmaffeo2 committed
1106
        ## TODO: decide whether to remove bead_model argument
1107
        ##       (currently unused)
cmaffeo2's avatar
cmaffeo2 committed
1108

1109
        tmp_children = []
1110

1111
1112
1113
1114
1115
        if one_bead_per_monomer:
            last = None
            for i in range(self.num_nt):
                s = self.nt_pos_to_contour(i)
                b = self._generate_one_bead(s,1)
1116

1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
                if last is not None:
                    last.make_intrahelical_neighbor(b)
                last = b
                tmp_children.append(b)
        else:
            ## First find points between-which beads must be generated
            # 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]
            # if self.name == "S001":
            #     pdb.set_trace()
cmaffeo2's avatar
cmaffeo2 committed
1128

1129
1130
1131
1132
            # pdb.set_trace()
            existing_beads0 = { (l.particle, l.particle.get_contour_position(self,l.address))
                                for l in self.locations if l.particle is not None }
            existing_beads = sorted( list(existing_beads0), key=lambda bc: bc[1] )
1133

1134
            # if self.num_nt == 1 and all([l.particle is not None for l in self.locations]):
cmaffeo2's avatar
cmaffeo2 committed
1135
            #     pdb.set_trace()
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
            #     return

            for b,c in existing_beads:
                assert(b.parent is not None)

            ## Add ends if they don't exist yet
            ## TODOTODO: test 1 nt segments?
            if len(existing_beads) == 0 or (existing_beads[0][0].get_nt_position(self,0) >= 0.5 and 'is_circular_bead' not in existing_beads[0][0].__dict__):
                # if len(existing_beads) > 0:            
                #     assert(existing_beads[0].get_nt_position(self) >= 0.5)
                c = self.nt_pos_to_contour(0)
                if self.num_nt == 1: c -= 0.4
                b = self._generate_one_bead(c, 0)
                existing_beads = [(b,0)] + existing_beads

            if (existing_beads[-1][0].get_nt_position(self,1)-(self.num_nt-1) < -0.5 and 'is_circular_bead' not in existing_beads[0][0].__dict__) or len(existing_beads)==1:
                c = self.nt_pos_to_contour(self.num_nt-1)
                if self.num_nt == 1: c += 0.4
                b = self._generate_one_bead(c, 0)
                existing_beads.append( (b,1) )
            assert(len(existing_beads) > 1)

            ## Walk through existing_beads, add beads between
            last = None

            for I in range(len(existing_beads)-1):
                eb1,eb2 = [existing_beads[i][0] for i in (I,I+1)]
                ec1,ec2 = [existing_beads[i][1] for i in (I,I+1)]
                assert( (eb1,ec1) is not (eb2,ec2) )

                # 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]))
                e_ds = ec2-ec1
                num_beads = self._get_num_beads( e_ds, max_basepairs_per_bead, max_nucleotides_per_bead )

                ## Ensure there is a ssDNA bead between dsDNA beads
                if num_beads == 0 and isinstance(self,SingleStrandedSegment) and isinstance(eb1.parent,DoubleStrandedSegment) and isinstance(eb2.parent,DoubleStrandedSegment):
                    num_beads = 1
                ## TODO similarly ensure there is a dsDNA bead between ssDNA beads

                ds = e_ds / (num_beads+1)
                nts = ds*self.num_nt
                eb1.num_nt += 0.5*nts
                eb2.num_nt += 0.5*nts

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

                s0 = ec1
                if last is not None:
                    last.make_intrahelical_neighbor(eb1)
                last = eb1
                for j in range(num_beads):
                    s = ds*(j+1) + s0
                    # if self.name in ("51-2","51-3"):
                    # if self.name in ("31-2",):
                    #     print(" adding bead at {}".format(s))
                    b = self._generate_one_bead(s,nts)
cmaffeo2's avatar
cmaffeo2 committed
1198

1199
1200
1201
                    last.make_intrahelical_neighbor(b)
                    last = b
                    tmp_children.append(b)
1202

1203
1204
1205
1206
            last.make_intrahelical_neighbor(eb2)

            if eb2.parent == self:
                tmp_children.append(eb2)
1207

cmaffeo2's avatar
cmaffeo2 committed
1208
        # if self.name in ("31-2",):
1209
        #     pdb.set_trace()
1210
        self._rebuild_children(tmp_children)
1211

cmaffeo2's avatar
cmaffeo2 committed
1212
1213
1214
        for callback in self._generate_bead_callbacks:
            callback(self)

1215
1216
1217
1218
1219
1220
1221
1222
1223
    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 """

1224
    def __init__(self, name, num_bp, start_position = np.array((0,0,0)),
1225
1226
                 end_position = None, 
                 segment_model = None,
cmaffeo2's avatar
cmaffeo2 committed
1227
1228
                 local_twist = False,
                 num_turns = None,
cmaffeo2's avatar
cmaffeo2 committed
1229
                 start_orientation = None,
cmaffeo2's avatar
cmaffeo2 committed
1230
1231
                 twist_persistence_length = 90,
                 **kwargs):
cmaffeo2's avatar
cmaffeo2 committed
1232
1233
1234
        
        self.helical_rise = 10.44
        self.distance_per_nt = 3.4
1235
        Segment.__init__(self, name, num_bp,
1236
1237
                         start_position,
                         end_position, 
cmaffeo2's avatar
cmaffeo2 committed
1238
1239
                         segment_model,
                         **kwargs)
1240
        self.num_bp = self.num_nt
1241

cmaffeo2's avatar
cmaffeo2 committed
1242
1243
        self.local_twist = local_twist
        if num_turns is None:
1244
1245
            num_turns = float(num_bp) / self.helical_rise
        self.twist_per_nt = float(360 * num_turns) / num_bp
cmaffeo2's avatar
cmaffeo2 committed
1246
1247

        if start_orientation is None:
1248
            start_orientation = np.eye(3) # np.array(((1,0,0),(0,1,0),(0,0,1)))
cmaffeo2's avatar
cmaffeo2 committed
1249
        self.start_orientation = start_orientation
cmaffeo2's avatar
cmaffeo2 committed
1250
        self.twist_persistence_length = twist_persistence_length
cmaffeo2's avatar
cmaffeo2 committed
1251

1252
1253
        self.nicks = []

1254
        self.start = self.start5 = Location( self, address=0, type_= "end5" )
1255
        self.start3 = Location( self, address=0, type_ = "end3", on_fwd_strand=False )
1256

1257
1258
        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
1259
1260
        # for l in (self.start5,self.start3,self.end3,self.end5):
        #     self.locations.append(l)
1261

1262
1263
        ## TODO: initialize sensible spline for orientation

1264
    ## Convenience methods
1265
    ## TODO: add errors if unrealistic connections are made
1266
    ## TODO: make connections automatically between unconnected strands
1267
    def connect_start5(self, end3, type_="intrahelical", force_connection=False):
1268
1269
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
1270
1271
        self._connect_ends( self.start5, end3, type_, force_connection = force_connection )
    def connect_start3(self, end5, type_="intrahelical", force_connection=False):
1272
        if isinstance(end5, SingleStrandedSegment):
1273
            end5 = end5.start5
1274
1275
        self._connect_ends( self.start3, end5, type_, force_connection = force_connection )
    def connect_end3(self, end5, type_="intrahelical", force_connection=False):
1276
        if isinstance(end5, SingleStrandedSegment):
1277
            end5 = end5.start5
1278
1279
        self._connect_ends( self.end3, end5, type_, force_connection = force_connection )
    def connect_end5(self, end3, type_="intrahelical", force_connection=False):
1280
1281
        if isinstance(end3, SingleStrandedSegment):
            end3 = end3.end3
1282
        self._connect_ends( self.end5, end3, type_, force_connection = force_connection )
1283

1284
    def add_crossover(self, nt, other, other_nt, strands_fwd=(True,False), nt_on_5prime=True, type_="crossover"):
cmaffeo2's avatar
cmaffeo2 committed
1285
1286
1287
1288
        """ Add a crossover between two helices """
        ## Validate other, nt, other_nt
        ##   TODO

1289
        if isinstance(other,SingleStrandedSegment):
cmaffeo2's avatar
cmaffeo2 committed
1290
            other.add_crossover(other_nt, self, nt, strands_fwd[::-1], not nt_on_5prime)
1291
        else:
1292

1293
1294
1295
            ## Create locations, connections and add to segments
            c = self.nt_pos_to_contour(nt)
            assert(c >= 0 and c <= 1)
1296

1297
1298
1299
1300
1301
1302

            if nt == 0 and not strands_fwd[0]:
                loc = self.start3
            elif nt == self.num_nt-1 and strands_fwd[0]:
                loc = self.end3
            else:
1303
                loc = self.get_location_at(c, strands_fwd[0], new_type='crossover')
1304
1305

            c = other.nt_pos_to_contour(other_nt)
cmaffeo2's avatar
cmaffeo2 committed
1306
            # TODOTODO: may need to subtract or add a little depending on 3prime/5prime
1307
            assert(c >= 0 and c <= 1)
1308
1309
1310
1311
1312
1313

            if other_nt == 0 and strands_fwd[1]:
                other_loc = other.start5
            elif other_nt == other.num_nt-1 and not strands_fwd[1]:
                other_loc = other.end5
            else:
1314
                other_loc = other.get_location_at(c, strands_fwd[1], new_type='crossover')
1315

1316
            self._connect(other, Connection( loc, other_loc, type_=type_ ))
cmaffeo2's avatar
cmaffeo2 committed
1317
1318
1319
1320
1321
1322
            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
1323

1324
    def _get_num_beads(self, contour, max_basepairs_per_bead, max_nucleotides_per_bead):
1325
1326
        # return int(contour*self.num_nt // max_basepairs_per_bead)
        return int(contour*(self.num_nt**2/(self.num_nt+1)) // max_basepairs_per_bead)
cmaffeo2's avatar
cmaffeo2 committed
1327

1328
1329
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
1330
        if self.local_twist:
1331
            orientation = self.contour_to_orientation(contour_position)
cmaffeo2's avatar
cmaffeo2 committed
1332
1333
1334
            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
1335
            opos = pos + orientation.dot( np.array((Segment.orientation_bond.r0,0,0)) )
cmaffeo2's avatar
cmaffeo2 committed
1336
1337
1338
1339
1340
            # 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,
1341
                                 num_nt=nts, parent=self )
1342
            bead = SegmentParticle( Segment.dsDNA_particle, pos, name="DNA",
1343
                                    num_nt=nts, parent=self, 
1344
1345
                                    orientation_bead=o,
                                    contour_position=contour_position )
cmaffeo2's avatar
cmaffeo2 committed
1346
1347

        else:
1348
            bead = SegmentParticle( Segment.dsDNA_particle, pos, name="DNA",
1349
                                    num_nt=nts, parent=self,
1350
1351
                                    contour_position=contour_position )
        self._add_bead(bead)
cmaffeo2's avatar
cmaffeo2 committed
1352
        return bead
1353
1354
1355
1356
1357
1358

class SingleStrandedSegment(Segment):

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

1359
    def __init__(self, name, num_nt, start_position = None,
1360
                 end_position = None, 
cmaffeo2's avatar
cmaffeo2 committed
1361
1362
                 segment_model = None,
                 **kwargs):
1363

cmaffeo2's avatar
cmaffeo2 committed
1364
        if start_position is None: start_position = np.array((0,0,0))
1365
        self.distance_per_nt = 5
1366
        Segment.__init__(self, name, num_nt, 
1367
1368
                         start_position,
                         end_position, 
cmaffeo2's avatar
cmaffeo2 committed
1369
1370
                         segment_model,
                         **kwargs)
1371

1372
        self.start = self.start5 = Location( self, address=0, type_= "end5" ) # TODO change type_?
1373
        self.end = self.end3 = Location( self, address=1, type_ = "end3" )
cmaffeo2's avatar
cmaffeo2 committed
1374
1375
        # for l in (self.start5,self.end3):
        #     self.locations.append(l)
1376

1377
    def connect_end3(self, end5, force_connection=False):
cmaffeo2's avatar
cmaffeo2 committed
1378
        self._connect_end( end5,  _5_to_3 = True, force_connection = force_connection )
1379

1380
    def connect_start5(self, end3, force_connection=False):
cmaffeo2's avatar
cmaffeo2 committed
1381
        self._connect_end( end3,  _5_to_3 = False, force_connection = force_connection )
1382
1383
1384
1385

    def _connect_end(self, other, _5_to_3, force_connection):
        assert( isinstance(other, Location) )
        if _5_to_3 == True:
1386
1387
1388
1389
            seg1 = self
            seg2 = other.container
            end1 = self.end3
            end2 = other
cmaffeo2's avatar
cmaffeo2 committed
1390
1391
1392
            assert(other.type_ != "end3")
            # if (other.type_ is not "end5"):
            #     print("WARNING: code does not prevent connecting 3prime to 3prime, etc")
1393
        else:
1394
1395
1396
            seg1 = other.container
            seg2 = self
            end1 = other
1397
            end2 = self.start
cmaffeo2's avatar
cmaffeo2 committed
1398
1399
1400
            assert(other.type_ != "end5")
            # if (other.type_ is not "end3"):
            #     print("WARNING: code does not prevent connecting 3prime to 3prime, etc")
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412

        ## Remove other connections involving these points
        if end1.connection is not None:
            print("WARNING: reconnecting {}".format(end1))
            end1.connection.delete()
        if end2.connection is not None:
            print("WARNING: reconnecting {}".format(end2))
            end2.connection.delete()

        conn = Connection( end1, end2, type_="intrahelical" )
        seg1._connect( seg2, conn, in_3prime_direction=True )

1413

1414
    def add_crossover(self, nt, other, other_nt, strands_fwd=(True,False), nt_on_5prime=True, type_='sscrossover'):
1415
        """ Add a crossover between two helices """
1416
1417
        ## TODO Validate other, nt, other_nt

1418
1419
        assert(nt < self.num_nt)
        assert(other_nt < other.num_nt)
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
        if nt_on_5prime:
            if nt == self.num_nt-1 and other_nt in (0,other.num_nt-1):
                other_end = None
                if strands_fwd[1] and other_nt == 0:
                    other_end = other.start5
                elif (not strands_fwd[1]) and other_nt == other.num_nt-1:
                    other_end = other.end5
                if other_end is not None:
                    self.connect_end3( other_end )
                    return
        else:
            if nt == 0 and other_nt in (0,other.num_nt-1):
                other_end = None
                if strands_fwd[1] and other_nt == other.num_nt-1:
                    other_end = other.end3
                elif (not strands_fwd[1]) and other_nt == 0:
                    other_end = other.start3
                if other_end is not None:
                    self.connect_start5( other_end )
                    return
1440

cmaffeo2's avatar
cmaffeo2 committed
1441
1442
1443
1444
        # 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))
1445
        # assert(np.isclose(nt,0) or np.isclose(nt,self.num_nt-1))
1446
        if nt == 0 and (self.num_nt > 1 or not nt_on_5prime):
cmaffeo2's avatar
cmaffeo2 committed
1447
            c1 = 0
1448
        elif nt == self.num_nt-1:
cmaffeo2's avatar
cmaffeo2 committed
1449
1450
1451
            c1 = 1
        else:
            raise Exception("Crossovers can only be at the ends of an ssDNA segment")
1452
        loc = self.get_location_at(c1, True, new_type='crossover')
1453

cmaffeo2's avatar
cmaffeo2 committed
1454
1455
        if other_nt == 0:
            c2 = 0
1456
        elif other_nt == other.num_nt-1:
cmaffeo2's avatar
cmaffeo2 committed
1457
1458
1459
1460
            c2 = 1
        else:
            c2 = other.nt_pos_to_contour(other_nt)

1461
1462
        if isinstance(other,SingleStrandedSegment):
            ## Ensure connections occur at opposing ends
1463
            assert(np.isclose(other_nt,0) or np.isclose(other_nt,self.num_nt-1))
1464
            other_loc = other.get_location_at( c2, True, new_type='crossover')
1465
1466
            # if ("22-2" in (self.name, other.name)):
            #     pdb.set_trace()
cmaffeo2's avatar
cmaffeo2 committed
1467
            if nt_on_5prime:
1468
1469
1470
1471
1472
1473
                self.connect_end3( other_loc )
            else:
                other.connect_end3( self )

        else:
            assert(c2 >= 0 and c2 <= 1)
1474
            other_loc = other.get_location_at( c2, strands_fwd[1], new_type='crossover')
cmaffeo2's avatar
cmaffeo2 committed
1475
            if nt_on_5prime:
1476
1477
1478
                self._connect(other, Connection( loc, other_loc, type_="sscrossover" ), in_3prime_direction=True )
            else:
                other._connect(self, Connection( other_loc, loc, type_="sscrossover" ), in_3prime_direction=True )
1479

1480
    def _get_num_beads(self, contour, max_basepairs_per_bead, max_nucleotides_per_bead):
1481
1482
        return int(contour*(self.num_nt**2/(self.num_nt+1)) // max_basepairs_per_bead)
        # return int(contour*self.num_nt // max_nucleotides_per_bead)
cmaffeo2's avatar
cmaffeo2 committed
1483

1484
1485
    def _generate_one_bead(self, contour_position, nts):
        pos = self.contour_to_position(contour_position)
1486
1487
        b = SegmentParticle( Segment.ssDNA_particle, pos, 
                             name="NAS",
1488
                             num_nt=nts, parent=self,
1489
1490
1491
                             contour_position=contour_position )
        self._add_bead(b)
        return b
1492
    
cmaffeo2's avatar
cmaffeo2 committed
1493
class StrandInSegment(Group):
1494
    """ Represents a piece of an ssDNA strand within a segment """
cmaffeo2's avatar
cmaffeo2 committed
1495
1496
    
    def __init__(self, segment, start, end, is_fwd):
1497
        """ start/end should be provided expressed in nt coordinates, is_fwd tuples """
cmaffeo2's avatar
cmaffeo2 committed
1498
        Group.__init__(self)
1499
        self.num_nt = 0
1500
        # self.sequence = []
cmaffeo2's avatar
cmaffeo2 committed
1501
1502
1503
1504
1505
        self.segment = segment
        self.start = start
        self.end = end
        self.is_fwd = is_fwd

1506
        nts = np.abs(end-start)+1
1507
1508
        self.num_nt = int(round(nts))
        assert( np.isclose(self.num_nt,nts) )
1509
        segment._add_strand_piece(self)
1510
1511
1512
1513
1514
1515
1516
1517
    
    def _nucleotide_ids(self):
        nt0 = self.start # seg.contour_to_nt_pos(self.start)
        assert( np.abs(nt0 - round(nt0)) < 1e-5 )
        nt0 = int(round(nt0))
        assert( (self.end-self.start) >= 0 or not self.is_fwd )

        direction = (2*self.is_fwd-1)
1518
        return range(nt0,nt0 + direction*self.num_nt, direction)
1519
1520
1521

    def get_sequence(self):
        """ return 5-to-3 """
1522
        # TODOTODO test
1523
1524
        seg = self.segment
        if self.is_fwd:
1525
            return [seg.sequence[nt] for nt in self._nucleotide_ids()]
1526
        else:
1527
            return [seqComplement[seg.sequence[nt]] for nt in self._nucleotide_ids()]
1528
1529
1530
    
    def get_contour_points(self):
        c0,c1 = [self.segment.nt_pos_to_contour(p) for p in (self.start,self.end)]
1531
        return np.linspace(c0,c1,self.num_nt)
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543

    def get_nucleotide(self, idx):
        """ idx expressed as nt coordinate within segment """

        lo,hi = sorted((self.start,self.end))
        if self.is_fwd:
            idx_in_strand = idx - lo
        else:
            idx_in_strand = hi - idx
        assert( np.isclose( idx_in_strand , int(round(idx_in_strand)) ) )
        assert(idx_in_strand >= 0)
        return self.children[int(round(idx_in_strand))]
1544
1545
    def __repr__(self):
        return "<StrandInSegment {}{}[{:.2f}-{:.2f},{:d}]>".format( self.parent.segname, self.segment.name, self.start, self.end, self.is_fwd)
1546

1547
            
cmaffeo2's avatar
cmaffeo2 committed
1548
class Strand(Group):
1549
    """ Represents an entire ssDNA strand from 5' to 3' as it routes through segments """
cmaffeo2's avatar
cmaffeo2 committed
1550
    def __init__(self, segname = None, is_circular = False):
cmaffeo2's avatar
cmaffeo2 committed
1551
        Group.__init__(self)
1552
        self.num_nt = 0
cmaffeo2's avatar
cmaffeo2 committed
1553
        self.children = self.strand_segments = []
1554
        self.oxdna_nt = []
1555
        self.segname = segname
cmaffeo2's avatar
cmaffeo2 committed
1556
        self.is_circular = is_circular
1557
        self.debug = False
1558

1559
1560
1561
    def __repr__(self):
        return "<Strand {}({})>".format( self.segname, self.num_nt )

1562
    ## TODO disambiguate names of functions
cmaffeo2's avatar
cmaffeo2 committed
1563
    def add_dna(self, segment, start, end, is_fwd):
1564
1565
        """ start/end are given as nt """
        if np.abs(start-end) <= 0.9:
1566
1567
            if self.debug:
                print( "WARNING: segment constructed with a very small number of nts ({})".format(np.abs(start-end)) )
1568
1569
            # import pdb
            # pdb.set_trace()
1570
1571
        for s in self.strand_segments:
            if s.segment == segment and s.is_fwd == is_fwd:
1572
1573
1574
                # assert( s.start not in (start,end) )
                # assert( s.end not in (start,end) )
                if s.start in (start,end) or s.end in (start,end):
cmaffeo2's avatar
cmaffeo2 committed
1575
                    raise CircularDnaError("Found circular DNA")
1576

cmaffeo2's avatar
cmaffeo2 committed
1577
        s = StrandInSegment( segment, start, end, is_fwd )
1578
        self.add( s )
1579
        self.num_nt += s.num_nt
cmaffeo2's avatar
cmaffeo2 committed
1580

New Tbgl User's avatar
New Tbgl User committed
1581
    def set_sequence(self,sequence): # , set_complement=True):
1582
        ## validate input
1583
        assert( len(sequence) >= self.num_nt )
New Tbgl User's avatar
New Tbgl User committed
1584
        assert( np.all( [i in ('A','T','C','G') for i in sequence] ) )
1585
1586
        
        seq_idx = 0
1587
1588
1589
        ## set sequence on each segment
        for s in self.children:
            seg = s.segment
1590
            if seg.sequence is None:
1591
                seg.sequence = [None for i in range(seg.num_nt)]
1592

New Tbgl User's avatar
New Tbgl User committed
1593
1594
1595
1596
            if s.is_fwd:
                for nt in s._nucleotide_ids():
                    seg.sequence[nt] = sequence[seq_idx]
                    seq_idx += 1
1597
            else:
New Tbgl User's avatar
New Tbgl User committed
1598
1599
1600
                for nt in s._nucleotide_ids():
                    seg.sequence[nt] = seqComplement[sequence[seq_idx]]
                    seq_idx += 1
1601
1602
1603
1604
1605
1606

    # def get_sequence(self):
    #     sequence = []
    #     for ss in self.strand_segments:
    #         sequence.extend( ss.get_sequence() )

1607
    #     assert( len(sequence) >= self.num_nt )
1608
1609
1610
    #     ret = ["5"+sequence[0]] +\
    #           sequence[1:-1] +\
    #           [sequence[-1]+"3"]
1611
    #     assert( len(ret) == self.num_nt )
1612
1613
    #     return ret

cmaffeo2's avatar
cmaffeo2 committed
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
    def link_nucleotides(self, nt5, nt3):
        parent = nt5.parent if nt5.parent is nt3.parent else self
        o3,c3,c4,c2,h3 = [nt5.atoms_by_name[n]
                          for n in ("O3'","C3'","C4'","C2'","H3'")]
        p,o5,o1,o2,c5 = [nt3.atoms_by_name[n]
                         for n in ("P","O5'","O1P","O2P","C5'")]
        parent.add_bond( o3, p, None )
        parent.add_angle( c3, o3, p, None )
        for x in (o5,o1,o2):
            parent.add_angle( o3, p, x, None )
            parent.add_dihedral(c3, o3, p, x, None )
        for x in (c4,c2,h3):
            parent.add_dihedral(x, c3, o3, p, None )
        parent.add_dihedral(o3, p, o5, c5, None)

1629
    def generate_atomic_model(self, scale, first_atomic_index):
cmaffeo2's avatar
cmaffeo2 committed
1630
1631
        last = None
        resid = 1
cmaffeo2's avatar
cmaffeo2 committed
1632
        ## TODO relabel "strand_segment"
cmaffeo2's avatar
cmaffeo2 committed
1633
        strand_segment_count = 0
cmaffeo2's avatar
cmaffeo2 committed
1634
        for s in self.strand_segments:
cmaffeo2's avatar
cmaffeo2 committed
1635
            strand_segment_count += 1
cmaffeo2's avatar
cmaffeo2 committed
1636
            seg = s.segment
1637
            contour = s.get_contour_points()
1638
1639
            # if s.end == s.start:
            #     pdb.set_trace()
cmaffeo2's avatar
cmaffeo2 committed
1640
            # assert(s.end != s.start)
1641
            assert( s.num_nt == 1 or (np.linalg.norm( seg.contour_to_position(contour[-1]) - seg.contour_to_position(contour[0]) ) > 0.1) )
1642
            nucleotide_count = 0
1643
            for c,seq in zip(contour,s.get_sequence()):
1644
                nucleotide_count += 1
cmaffeo2's avatar
cmaffeo2 committed
1645
                if last is None and not self.is_circular:
cmaffeo2's avatar
cmaffeo2 committed
1646
                    seq = "5"+seq
1647
                if strand_segment_count == len(s.strand_segments) and nucleotide_count == s.num_nt and not self.is_circular:
cmaffeo2's avatar
cmaffeo2 committed
1648
1649
                    seq = seq+"3"

cmaffeo2's avatar
cmaffeo2 committed
1650
1651
                nt = seg._generate_atomic_nucleotide( c, s.is_fwd, seq, scale, s )

cmaffeo2's avatar
cmaffeo2 committed
1652
1653
                ## Join last basepairs
                if last is not None:
cmaffeo2's avatar
cmaffeo2 committed
1654
1655
                    self.link_nucleotides(last,nt)

cmaffeo2's avatar
cmaffeo2 committed
1656
1657
1658
                nt.__dict__['resid'] = resid
                resid += 1
                last = nt
1659
1660
                nt._first_atomic_index = first_atomic_index
                first_atomic_index += len(nt.children)
cmaffeo2's avatar
cmaffeo2 committed
1661
1662
1663
1664

        if self.is_circular:
            self.link_nucleotides(last,self.strand_segments[0].children[0])

1665
        return first_atomic_index
1666

1667
    def generate_atomic_model(self, scale, first_atomic_index):
1668
1669
        last = None
        resid = 1
1670
1671
        ## TODO relabel "strand_segment"
        strand_segment_count = 0
1672
        for s in self.strand_segments:
1673
            strand_segment_count += 1
1674
1675
            seg = s.segment
            contour = s.get_contour_points()
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
            # if s.end == s.start:
            #     pdb.set_trace()
            # assert(s.end != s.start)
            assert( s.num_nt == 1 or (np.linalg.norm( seg.contour_to_position(contour[-1]) - seg.contour_to_position(contour[0]) ) > 0.1) )
            nucleotide_count = 0
            for c,seq in zip(contour,s.get_sequence()):
                nucleotide_count += 1
                if last is None and not self.is_circular:
                    seq = "5"+seq
                if strand_segment_count == len(s.strand_segments) and nucleotide_count == s.num_nt and not self.is_circular:
                    seq = seq+"3"

                nt = seg._generate_atomic_nucleotide( c, s.is_fwd, seq, scale, s )

                ## Join last basepairs
                if last is not None:
                    self.link_nucleotides(last,nt)

                nt.__dict__['resid'] = resid
                resid += 1
                last = nt
                nt._first_atomic_index = first_atomic_index
                first_atomic_index += len(nt.children)

        if self.is_circular:
            self.link_nucleotides(last,self.strand_segments[0].children[0])

        return first_atomic_index

    def generate_oxdna_model(self):
        for s in self.strand_segments:
            seg = s.segment
            contour = s.get_contour_points()
            assert( s.num_nt == 1 or (np.linalg.norm( seg.contour_to_position(contour[-1]) - seg.contour_to_position(contour[0]) ) > 0.1) )
            for c,seq in zip(contour,s.get_sequence()):
                nt = seg._generate_oxdna_nucleotide( c, s.is_fwd, seq )
                self.oxdna_nt.append(nt)
1713

1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
    def get_sequence(self):
        return ''.join([''.join(s.get_sequence()) for s in self.strand_segments])

    def get_nt_positions_orientations(self, bp_offset = (5,0,0)):
        ## TODO: remove duplicate code
        rs = []
        os = []
        for s in self.strand_segments:
            seg = s.segment
            contour = s.get_contour_points()
            assert( s.num_nt == 1 or (np.linalg.norm( seg.contour_to_position(contour[-1]) - seg.contour_to_position(contour[0]) ) > 0.1) )

            DefaultOrientation = rotationAboutAxis([0,0,1], 90)
            if s.is_fwd:
                DefaultOrientation = rotationAboutAxis([1,0,0], 180).dot(DefaultOrientation)

            for c,seq in zip(contour,s.get_sequence()):
                r = seg.contour_to_position(c)
                o = seg.contour_to_orientation(c)
                o = o.dot(DefaultOrientation)

                if isinstance(seg, SingleStrandedSegment):
                    r = r
                else:
                   r = r - o.dot(np.array(bp_offset))

                rs.append(r)
                os.append(quaternion_from_matrix(o))
        return np.array(rs), np.array(os)
1743
1744


1745
class SegmentModel(ArbdModel):
cmaffeo2's avatar
cmaffeo2 committed
1746
    def __init__(self, segments=[], local_twist=True, escapable_twist=True,
cmaffeo2's avatar
cmaffeo2 committed
1747
1748
                 max_basepairs_per_bead=7,
                 max_nucleotides_per_bead=4,
1749
                 origin = None,
1750
                 dimensions=(5000,5000,5000), temperature=291,
1751
1752
                 timestep=50e-6, cutoff=50, 
                 decompPeriod=10000, pairlistDistance=None, 
1753
                 nonbondedResolution=0, DEBUG=0,
1754
                 integrator='Brown',
1755
                 debye_length = None,
1756
                 hj_equilibrium_angle = 0,
1757
                 extra_bd_file_lines = "",
1758
    ):
1759
1760
        self.DEBUG = DEBUG
        if DEBUG > 0: print("Building ARBD Model")
1761
        ArbdModel.__init__(self,segments,
cmaffeo2's avatar
cmaffeo2 committed
1762
1763
1764
1765
                           origin, dimensions,
                           temperature, timestep, integrator, cutoff,
                           decompPeriod, pairlist_distance=None,
                           nonbonded_resolution = 0,
1766
1767
                           extra_bd_file_lines = extra_bd_file_lines
        )
1768

1769

cmaffeo2's avatar
cmaffeo2 committed
1770
1771
        # self.max_basepairs_per_bead = max_basepairs_per_bead     # dsDNA
        # self.max_nucleotides_per_bead = max_nucleotides_per_bead # ssDNA
1772
1773
        self.children = self.segments = segments

1774
1775
1776
1777
        if (hj_equilibrium_angle > 180) or (hj_equilibrium_angle < -180):
            raise ValueError("Holliday junction equilibrium dihedral angle must be in the range from -180 to 180 degrees!")
        self.hj_equilibrium_angle = hj_equilibrium_angle

1778
1779
        self._generate_bead_callbacks = []

cmaffeo2's avatar
cmaffeo2 committed
1780
        self._bonded_potential = dict() # cache for bonded potentials
1781
        self._generate_strands()
1782
        self.grid_potentials = []
cmaffeo2's avatar
cmaffeo2 committed
1783
        self._generate_bead_model( max_basepairs_per_bead, max_nucleotides_per_bead, local_twist, escapable_twist)
1784

1785
1786
        if debye_length is not None:
            nbDnaScheme.debye_length = debye_length
1787
        self.debye_length = debye_length
1788

1789
        self.useNonbondedScheme( nbDnaScheme )
1790
        self.useTclForces = False
1791

1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
    def set_debye_length(self, value):
        if value <= 0:
            raise ValueError("The Debye length must be positive")
        for s,tA,tB in self.nbSchemes:
            try:
                s.debye_length = value
            except:
                ...
        self.debye_length = value

cmaffeo2's avatar
cmaffeo2 committed
1802
    def get_connections(self,type_=None,exclude=()):
1803
1804
1805
1806
        """ Find all connections in model, without double-counting """
        added=set()
        ret=[]
        for s in self.segments:
cmaffeo2's avatar
cmaffeo2 committed
1807
            items = [e for e in s.get_connections_and_locations(type_,exclude=exclude) if e[0] not in added]
1808
            added.update([e[0] for e in items])
cmaffeo2's avatar
cmaffeo2 committed
1809
            ret.extend( list(sorted(items,key=lambda x: x[1].address)) )
1810
        return ret
1811
    
cmaffeo2's avatar
cmaffeo2 committed
1812
    def _recursively_get_beads_within_bonds(self,b1,bonds,done=()):
1813
        ret = []
1814
1815
1816
1817
1818
1819
1820
1821
1822
        done = list(done)
        done.append(b1)
        if bonds == 0:
            return [[]]

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

1825
1826
1827
1828
1829
    def _get_intrahelical_beads(self,num=2):
        ## TODO: add check that this is not called before adding intrahelical_neighbors in _generate_bead_model

        assert(num >= 2)

1830
1831
        ret = []
        for s in self.segments:
1832
1833
1834
1835
1836
            for b1 in s.beads:
                for bead_list in self._recursively_get_beads_within_bonds(b1, num-1):
                    assert(len(bead_list) == num-1)
                    if b1.idx < bead_list[-1].idx: # avoid double-counting
                        ret.append([b1]+bead_list)
1837
1838
        return ret

1839

1840
1841
    def _get_intrahelical_angle_beads(self):
        return self._get_intrahelical_beads(num=3)
1842

cmaffeo2's avatar
cmaffeo2 committed
1843
    def _get_potential(self, type_, kSpring, d, max_potential = None):
1844
        key = (type_, kSpring, d, max_potential)
1845
        if key not in self._bonded_potential:
cmaffeo2's avatar
cmaffeo2 committed
1846
            assert( kSpring >= 0 )
1847
            if type_ == "bond":
1848
                self._bonded_potential[key] = HarmonicBond(kSpring,d, rRange=(0,1200), max_potential=max_potential)
1849
1850
            elif type_ == "gbond":
                self._bonded_potential[key] = HarmonicBond(kSpring,d, rRange=(0,1200), max_potential=max_potential, correct_geometry=True, temperature=self.temperature)
1851
            elif type_ == "angle":
cmaffeo2's avatar
cmaffeo2 committed
1852
1853
                self._bonded_potential[key] = HarmonicAngle(kSpring,d, max_potential=max_potential)
                # , resolution = 1, maxForce=0.1)
1854
            elif type_ == "dihedral":
cmaffeo2's avatar
cmaffeo2 committed
1855
                self._bonded_potential[key] = HarmonicDihedral(kSpring,d, max_potential=max_potential)
1856
1857
1858
            else:
                raise Exception("Unhandled potential type '%s'" % type_)
        return self._bonded_potential[key]
1859
    def get_bond_potential(self, kSpring, d, correct_geometry=False):
1860
        assert( d > 0.2 )
1861
1862
        return self._get_potential("gbond" if correct_geometry else "bond",
                                   kSpring, d)
1863
1864
    def get_angle_potential(self, kSpring, d):
        return self._get_potential("angle", kSpring, d)
cmaffeo2's avatar
cmaffeo2 committed
1865
1866
1867
1868
    def get_dihedral_potential(self, kSpring, d, max_potential=None):
        while d > 180: d-=360
        while d < -180: d+=360
        return self._get_potential("dihedral", kSpring, d, max_potential)
1869

1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
    def _get_wlc_sk_bond_potential(self, d):
        ## https://aip.scitation.org/doi/full/10.1063/1.4968020
        type_='wclsk-bond'
        # lp = 2*5                # TODO place these parameters in a logical spot
        lp = 15                 # https://journals.aps.org/pre/abstract/10.1103/PhysRevE.86.021901
        # https://www.pnas.org/content/pnas/109/3/799.full.pdf
        lp = 12                 # selected semi-empirically to make ssDNA force extension match
        key = (type_, d, lp)
        assert( d > 0.2 )
        if key not in self._bonded_potential:
            kT = self.temperature * 0.0019872065 # kcal/mol
cmaffeo2's avatar
cmaffeo2 committed
1881
            self._bonded_potential[key] = WLCSKBond( d, lp, kT, rRange=(0,1200) )
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
        return self._bonded_potential[key]

    def _get_wlc_sk_angle_potential(self, d):
        ## https://aip.scitation.org/doi/full/10.1063/1.4968020
        type_='wclsk-angle'
        ## TODO move
        lp = 15                 # https://journals.aps.org/pre/abstract/10.1103/PhysRevE.86.021901
        # https://www.pnas.org/content/pnas/109/3/799.full.pdf
        lp = 12                 # selected semi-empirically to make ssDNA force extension match
        key = (type_, d, lp)
        assert( d > 0.2 )
        if key not in self._bonded_potential:
            kT = self.temperature * 0.0019872065 # kcal/mol
            self._bonded_potential[key] = WLCSKAngle( d, lp, kT )
        return self._bonded_potential[key]

1898

cmaffeo2's avatar
cmaffeo2 committed
1899
    def _getParent(self, *beads ):
1900
        if len(set([b.parent for b in beads])) == 1:
cmaffeo2's avatar
cmaffeo2 committed
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
            return beads[0].parent
        else:
            return self

    def _get_twist_spring_constant(self, sep):
        """ sep in nt """
        twist_persistence_length = 90  # set semi-arbitrarily as there is a large spread in literature
        ##   units "3e-19 erg cm/ 295 k K" "nm" =~ 73
        Lp = twist_persistence_length/0.34 

1911
        return twist_spring_from_lp(sep, Lp, self.temperature)
1912

cmaffeo2's avatar
cmaffeo2 committed
1913
    def extend(self, other, copy=True, include_strands=True):
1914
        assert( isinstance(other, SegmentModel) )
1915
1916
1917
1918
1919
1920
1921

        try:
            max_occupancy = max([s.occupancy for s in self.segments if 'occupancy' in s.__dict__])
            occupancy0 = 10**np.ceil(np.log10(max_occupancy+1))
        except:
            pass

1922
1923
        if copy:
            for s in other.segments:
1924
1925
1926
1927
1928
1929
                newseg = deepcopy(s)
                try:
                    newseg.occupancy = occupancy0+s.occupancy
                except:
                    pass
                self.segments.append(newseg)
1930
                newseg.parent = self
1931
1932
            if include_strands:
                for s in other.strands:
1933
1934
1935
                    newstrand = deepcopy(s)
                    self.strands.append(newstrand)
                    newstrand.parent = self
1936
1937
        else:
            for s in other.segments:
1938
1939
1940
1941
                try:
                    s.occupancy = occupancy0+s.occupancy
                except:
                    pass
1942
                self.segments.append(s)
1943
                s.parent = self
1944
1945
1946
            if include_strands:
                for s in other.strands:
                    self.strands.append(s)
1947
                    s.parent = self
1948
1949
1950
        self._clear_beads()

    def update(self, segment , copy=False):
1951
        assert( isinstance(segment, Segment) )
1952
1953
1954
1955
1956
        if copy:
            segment = deepcopy(segment)
        self.segments.append(segment)
        self._clear_beads()