segmentmodel.py 165 KB
Newer Older
2001

2002
2003
2004
2005
    def _update_segment_positions(self, bead_coordinates):
        print("WARNING: called deprecated command '_update_segment_positions; use 'update_splines' instead")
        return self.update_splines(bead_coordinates)

2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
    ## Operations on spline coordinates
    def translate(self, translation_vector, position_filter=None):
        for s in self.segments:
            s.translate(translation_vector, position_filter=position_filter)
        
    def rotate(self, rotation_matrix, about=None, position_filter=None):
        for s in self.segments:
            s.rotate(rotation_matrix, about=about, position_filter=position_filter)

    def get_center(self, include_ssdna=False):
        if include_ssdna:
            segments = self.segments
        else:
            segments = list(filter(lambda s: isinstance(s,DoubleStrandedSegment), 
                                   self.segments))
        centers = [s.get_center() for s in segments]
        weights = [s.num_nt*2 if isinstance(s,DoubleStrandedSegment) else s.num_nt for s in segments]
        # centers,weights = [np.array(a) for a in (centers,weights)]
        return np.average( centers, axis=0, weights=weights)

2026
    def update_splines(self, bead_coordinates):
2027
2028
        """ Set new function for each segments functions
        contour_to_position and contour_to_orientation """
cmaffeo2's avatar
cmaffeo2 committed
2029

2030
        for s in self.segments:
cmaffeo2's avatar
cmaffeo2 committed
2031
2032
2033
            # if s.name == "61-1":
            #     pdb.set_trace()

2034
            cabs = s.get_connections_and_locations("intrahelical")
2035
2036
            if isinstance(s,SingleStrandedSegment):
                cabs = cabs + [[c,A,B] for c,A,B in s.get_connections_and_locations("sscrossover") if A.address == 0 or A.address == 1]
cmaffeo2's avatar
cmaffeo2 committed
2037
2038
2039
            if np.any( [B.particle is None for c,A,B in cabs] ):
                print( "WARNING: none type found in connection, skipping" )
                cabs = [e for e in cabs if e[2].particle is not None]
2040

2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
            def get_beads_and_contour_positions(s):
                ret_list = []
                def zip_bead_contour(beads,address=None):
                    if isinstance(address,list):
                        assert(False)
                        for b,a in zip(beads,address):
                            if b is None: continue
                            try:
                                ret_list.append((b, b.get_contour_position(s,a)))
                            except:
                                ...                                
                    else:
                        for b in beads:
                            if b is None: continue
                            try:
                                ret_list.append((b, b.get_contour_position(s,address)))
                            except:
                                ...
                    return ret_list
                
                ## Add beads from segment s
                beads_contours = zip_bead_contour(s.beads)
                beads_contours.extend( zip_bead_contour([A.particle for c,A,B in cabs]) )
                beads = set([b for b,c in beads_contours])

                ## Add nearby beads
                for c,A,B in cabs:
                    ## TODOTODO test?
                    filter_fn = lambda x: x is not None and x not in beads
                    bs = list( filter( filter_fn, B.particle.intrahelical_neighbors ) )
                    beads_contours.extend( zip_bead_contour( bs, A.address ) )
2072
                    beads.update(bs)
2073
2074
2075
2076
                    for i in range(3):
                        bs = list( filter( filter_fn, [n for b in bs for n in b.intrahelical_neighbors] ) )
                        beads_contours.extend( zip_bead_contour( bs, A.address ) )
                        beads.update(bs)
2077

2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
                beads_contours = list(set(beads_contours))
                beads = list(beads)

                ## Skip beads that are None (some locations were not assigned a particle to avoid double-counting) 
                # beads = [b for b in beads if b is not None]
                assert( np.any([b is None for b,c in beads_contours]) == False )
                # beads = list(filter(lambda x: x[0] is not None, beads))

                if isinstance(s, DoubleStrandedSegment):
                    beads_contours = list(filter(lambda x: x[0].type_.name[0] == "D", beads_contours))
                return beads_contours
2089

2090
2091
            beads_contours = get_beads_and_contour_positions(s)
            contours = [c for b,c in beads_contours]
2092
2093
            contour_idx = np.array( np.array(contours)*s.num_nt * 10, dtype=np.int )
            contour_idx,ids1 = np.unique(contour_idx, return_index=True)
2094
2095
            beads_contours = [beads_contours[i] for i in ids1]

cmaffeo2's avatar
cmaffeo2 committed
2096
            ## TODO: keep closest beads beyond +-1.5 if there are fewer than 2 beads
2097
2098
2099
            tmp = []
            dist = 1
            while len(tmp) < 5 and dist < 3:
2100
                tmp = list(filter(lambda bc: np.abs(bc[1]-0.5) < dist, beads_contours))
2101
2102
2103
2104
2105
                dist += 0.1

            if len(tmp) <= 1:
                raise Exception("Failed to fit spline into segment {}".format(s))

2106
2107
            beads = [b for b,c in tmp]
            contours = [c for b,c in tmp]
2108
2109
            ids = [b.idx for b in beads]
            
cmaffeo2's avatar
cmaffeo2 committed
2110
2111
2112
            if len(beads) <= 1:
                pdb.set_trace()

cmaffeo2's avatar
cmaffeo2 committed
2113

2114
2115
            """ Get positions """
            positions = bead_coordinates[ids,:].T
2116

cmaffeo2's avatar
cmaffeo2 committed
2117
2118
2119
2120
2121
            # print("BEADS NOT IN {}:".format(s))
            # for b,c in filter(lambda z: z[0] not in s.beads, zip(beads,contours)):
            #     print("  {}[{:03f}]: {:0.3f}".format(b.parent.name, b.contour_position, c) )


2122
            tck, u = interpolate.splprep( positions, u=contours, s=0, k=1 )
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140

            # if len(beads) < 8:
            #     ret = interpolate.splprep( positions, u=contours, s=0, k=1, full_output=1 )
            #     tck = ret[0][0]
            #     if ret[2] > 0:
            #         pdb.set_trace()

            # else:
            #     try:
            #         ret = interpolate.splprep( positions, u=contours, s=0, k=3, full_output=1 )
            #         tck = ret[0][0]
            #         if ret[2] > 0:
            #             pdb.set_trace()
            #     except:
            #         ret = interpolate.splprep( positions, u=contours, s=0, k=1, full_output=1 )
            #         tck = ret[0][0]
            #         if ret[2] > 0:
            #             pdb.set_trace()
2141

2142
            s.position_spline_params = (tck,u)
2143

2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
            """ Get orientation """
            def get_orientation_vector(bead,tangent):
                if 'orientation_bead' in bead.__dict__:
                    o = bead.orientation_bead
                    oVec = bead_coordinates[o.idx,:] - bead_coordinates[bead.idx,:]
                    oVec = oVec - oVec.dot(tangent)*tangent
                    oVec = oVec/np.linalg.norm(oVec)
                else:
                    oVec = None
                return oVec

            def remove_tangential_projection(vector, tangent):
                """ Assume tangent is normalized """
                v = vector - vector.dot(tangent)*tangent
                return v/np.linalg.norm(v)

            def get_orientation_vector(bead,tangent):
                if 'orientation_bead' in bead.__dict__:
                    o = bead.orientation_bead
                    oVec = bead_coordinates[o.idx,:] - bead_coordinates[bead.idx,:]
                    oVec = remove_tangential_projection(oVec,tangent)
                else:
                    oVec = None
                return oVec


            def get_previous_idx_if_none(list_):
                previous = None
                result = []
                i = 0
                for e in list_:
                    if e is None:
                        result.append(previous)
                    else:
                        previous = i
                    i+=1
                return result
            def get_next_idx_if_none(list_):
                tmp = get_previous_idx_if_none(list_[::-1])[::-1]
                return [ len(list_)-1-idx if idx is not None else idx for idx in tmp ]

            def fill_in_orientation_vectors(contours,orientation_vectors,tangents):
                result = []
                last_idx = get_previous_idx_if_none( orientation_vectors )
                next_idx = get_next_idx_if_none( orientation_vectors )
                none_idx = 0
                for c,ov,t in zip(contours,orientation_vectors,tangents):
                    if ov is not None:
                        result.append(ov)
                    else:
                        p = last_idx[none_idx]
                        n = next_idx[none_idx]
                        none_idx += 1
                        if p is None:
                            if n is None:
                                ## Should be quite rare; give something random if it happens
                                print("WARNING: unable to interpolate orientation")
                                o = np.array((1,0,0))
                                result.append( remove_tangential_projection(o,t) )
                            else:
                                o = orientation_vectors[n]
                                result.append( remove_tangential_projection(o,t) )
                        else:
                            if n is None:
                                o = orientation_vectors[p]
                                result.append( remove_tangential_projection(o,t) )
                            else:
                                cp,cn = [contours[i] for i in (p,n)]
                                op,on = [orientation_vectors[i] for i in (p,n)]
                                if (cn-cp) > 1e-6:
                                    o = ((cn-c)*op+(c-cp)*on)/(cn-cp)
                                else:
                                    o = op+on
                                result.append( remove_tangential_projection(o,t) )
                return result

            tangents = s.contour_to_tangent(contours)
            orientation_vectors = [get_orientation_vector(b,t) for b,t in zip(beads,tangents)]
            if len(beads) > 3 and any([e is not None for e in orientation_vectors] ):
                orientation_vectors = fill_in_orientation_vectors(contours, orientation_vectors, tangents)

2225
                quats = []
cmaffeo2's avatar
cmaffeo2 committed
2226
                lastq = None
2227
2228
                for b,t,oVec in zip(beads,tangents,orientation_vectors):
                    y = np.cross(t,oVec)
cmaffeo2's avatar
cmaffeo2 committed
2229
                    assert( np.abs(np.linalg.norm(y) - 1) < 1e-2 )
2230
                    q = quaternion_from_matrix( np.array([oVec,y,t]).T)
2231
                    
cmaffeo2's avatar
cmaffeo2 committed
2232
2233
2234
2235
                    if lastq is not None:
                        if q.dot(lastq) < 0:
                            q = -q
                    quats.append( q )
2236
                    lastq = q
cmaffeo2's avatar
cmaffeo2 committed
2237

2238
                # pdb.set_trace()
2239
                quats = np.array(quats)
2240
2241
                # tck, u = interpolate.splprep( quats.T, u=contours, s=3, k=3 ) ;# cubic spline not as good
                tck, u = interpolate.splprep( quats.T, u=contours, s=0, k=1 )
2242
                s.quaternion_spline_params = (tck,u)
2243

2244
2245
2246
2247
    def get_segments_matching(self, pattern):
        """ Returns a list of all segments that match the regular expression 'pattern' """

        re_pattern = re.compile(pattern)
2248
        return [s for s in self.segments if re_pattern.match(s.name) is not None]
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291

    def get_crossovers_at_ends(self):
        """
        Return a list of all "regular" crossovers where both sides of
        the crossover exist at the start or end of a segment
        """

        ret_list = []
        for s in self.segments:
            for c in filter(lambda c: c.type_ == "crossover", s.connections):
                seg1 = c.A.container
                seg2 = c.B.container
                nt1 = c.A.get_nt_pos()
                nt2 = c.B.get_nt_pos()
                if (nt1 < 1 and nt2 > seg2.num_nt-2) or (nt2 < 1 and nt1 > seg1.num_nt-2):
                    if c not in ret_list:
                        ret_list.append(c)
        return ret_list

    def convert_crossover_to_intrahelical(self, crossover):
        c = crossover
        seg1 = c.A.container
        seg2 = c.B.container
        nt1 = c.A.get_nt_pos()
        nt2 = c.B.get_nt_pos()

        assert( c.A.on_fwd_strand == c.B.on_fwd_strand )
        on_fwd_strand = c.A.on_fwd_strand

        c.delete()

        if nt1 < 1 and nt2 > seg2.num_nt-2:
            if on_fwd_strand:
                seg2.connect_end3(seg1.start5)
            else:
                seg1.connect_start3(seg2.end5)

        elif nt2 < 1 and nt1 > seg1.num_nt-2:
            if on_fwd_strand:
                seg1.connect_end3(seg2.start5)
            else:
                seg2.connect_start3(seg1.end5)

2292
2293
2294
    def get_blunt_DNA_ends(self):
        return sum([s.get_blunt_DNA_ends for s in self.segments],[])

2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
    def convert_crossovers_at_ends_beyond_cutoff_to_intrahelical(self,cutoff):
        if cutoff < 0:
            raise ValueError("Cutoff should be positive!")

        for c in self.get_crossovers_at_ends():
            r1,r2 = [l.container.contour_to_position(l.address)
                     for l in (c.A,c.B)]
            if np.linalg.norm(r2-r1) > cutoff:
                self.convert_crossover_to_intrahelical(c)

2305
    def _generate_bead_model(self,
2306
                             max_basepairs_per_bead = 7,
cmaffeo2's avatar
cmaffeo2 committed
2307
                             max_nucleotides_per_bead = 4,
2308
2309
                             local_twist=False,
                             escapable_twist=True):
2310
2311
2312
2313
2314
        ## TODO: deprecate
        self.generate_bead_model( max_basepairs_per_bead = max_basepairs_per_bead,
                                  max_nucleotides_per_bead = max_nucleotides_per_bead,
                                  local_twist=local_twist,
                                  escapable_twist=escapable_twist)
2315

2316
    def generate_bead_model(self,
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
                            max_basepairs_per_bead = 7,
                            max_nucleotides_per_bead = 4,
                            local_twist=False,
                            escapable_twist=True,
                            sequence_dependent = False,
                            one_bead_per_monomer = False
                        ):
        if sequence_dependent and not local_twist:
            raise ValueError("Incompatible options: if sequence_dependent==True, then local_twist must be True")
        if sequence_dependent and not one_bead_per_monomer:
            raise ValueError("Incompatible options: if sequence_dependent==True, then one_bead_per_monomer must be True")
2328
2329
2330

        self.children = self.segments # is this okay?
        self.clear_beads()
2331

2332
        segments = self.segments
2333
2334
        for s in segments:
            s.local_twist = local_twist
2335
2336

        """ Simplify connections """
2337
2338
        # d_nt = dict()           # 
        # for s in segments:
2339
        #     d_nt[s] = 1.5/(s.num_nt-1)
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
        # for s in segments:
        #     ## replace consecutive crossovers with
        #     cl = sorted( s.get_connections_and_locations("crossover"), key=lambda x: x[1].address )
        #     last = None
        #     for entry in cl:
        #         c,A,B = entry
        #         if last is not None and \
        #            (A.address - last[1].address) < d_nt[s]:
        #             same_type = c.type_ == last[0].type_
        #             same_dest_seg = B.container == last[2].container
        #             if same_type and same_dest_seg:
        #                 if np.abs(B.address - last[2].address) < d_nt[B.container]:
        #                     ## combine
        #                     A.combine = last[1]
        #                     B.combine = last[2]

        #                     ...
        #         # if last is not None:
        #         #     s.bead_locations.append(last)
        #         ...
        #         last = entry
        # del d_nt
2362

2363
        """ Generate beads at intrahelical junctions """
2364
        if self.DEBUG: print( "Adding intrahelical beads at junctions" )
2365

2366
2367
2368
2369
        if not one_bead_per_monomer:
            ## Loop through all connections, generating beads at appropriate locations
            for c,A,B in self.get_connections("intrahelical"):
                s1,s2 = [l.container for l in (A,B)]
2370

2371
2372
                assert( A.particle is None )
                assert( B.particle is None )
2373

2374
2375
2376
2377
2378
                ## TODO: offload the work here to s1
                # TODOTODO
                a1,a2 = [l.address   for l in (A,B)]            
                for a in (a1,a2):
                    assert( np.isclose(a,0) or np.isclose(a,1) )
2379

2380
                ## TODO improve this for combinations of ssDNA and dsDNA (maybe a1/a2 should be calculated differently)
cmaffeo2's avatar
cmaffeo2 committed
2381

2382
2383
2384
2385
2386
2387
2388
2389
                def _get_nearest(seg,address):
                    b = seg.get_nearest_bead(address)
                    if b is not None:
                        assert( b.parent is seg )
                        """ if above assertion is true, no problem here """
                        if np.abs(b.get_nt_position(seg) - seg.contour_to_nt_pos(address)) > 0.5:
                            b = None
                    return b
2390

2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
                """ Search to see whether bead at location is already found """
                b = None
                if isinstance(s1,DoubleStrandedSegment):
                    b = _get_nearest(s1,a1)
                if b is None and isinstance(s2,DoubleStrandedSegment):
                    b = _get_nearest(s2,a2)

                if b is not None and b.parent not in (s1,s2):
                    b = None

                if b is None:
                    ## need to generate a bead
                    if isinstance(s2,DoubleStrandedSegment):
                        b = s2._generate_one_bead(a2,0)
                    else:
                        b = s1._generate_one_bead(a1,0)
                A.particle = B.particle = b
                b.is_intrahelical = True
                b.locations.extend([A,B])
                if s1 is s2:
                    b.is_circular_bead = True
2412

2413
            # pdb.set_trace()
cmaffeo2's avatar
cmaffeo2 committed
2414

2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
            """ Generate beads at other junctions """
            for c,A,B in self.get_connections(exclude="intrahelical"):
                s1,s2 = [l.container for l in (A,B)]
                if A.particle is not None and B.particle is not None:
                    continue
                # assert( A.particle is None )
                # assert( B.particle is None )

                ## TODO: offload the work here to s1/s2 (?)
                a1,a2 = [l.address   for l in (A,B)]

                def maybe_add_bead(location, seg, address, ):
                    if location.particle is None:
                        b = seg.get_nearest_bead(address)
                        try:
                            distance = seg.contour_to_nt_pos(np.abs(b.contour_position-address))
                            max_distance =  min(max_basepairs_per_bead, max_nucleotides_per_bead)*0.5
                            if "is_intrahelical" in b.__dict__:
                                max_distance = 0.5
                            if distance >= max_distance:
                                raise Exception("except")

                            ## combine beads
                            b.update_position( 0.5*(b.contour_position + address) ) # avg position
                        except:
                            b = seg._generate_one_bead(address,0)
                        location.particle = b
                        b.locations.append(location)

                maybe_add_bead(A,s1,a1)
                maybe_add_bead(B,s2,a2)

            """ Some tests """
            for c,A,B in self.get_connections("intrahelical"):
                for l in (A,B):
                    if l.particle is None: continue
                    assert( l.particle.parent is not None )
2452

2453
2454
2455
2456
            """ Generate beads in between """
            if self.DEBUG: print("Generating beads")
            for s in segments:
                s._generate_beads( self, max_basepairs_per_bead, max_nucleotides_per_bead )
2457

2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
        else:
            """ Generate beads in each segment """
            if self.DEBUG: print("Generating beads")
            for s in segments:
                s._generate_beads( self, max_basepairs_per_bead, max_nucleotides_per_bead, one_bead_per_monomer )

            """ Associate beads with junctions """
            # for c,A,B in self.get_connections("intrahelical"):
            for c,A,B in self.get_connections():
                for l in (A,B):
                    seg = l.container
                    a = l.address
                    b = seg.get_nearest_bead(a)

                    l.particle = b
                    b.locations.append(l)
                """
                b.is_intrahelical = True
                b.locations.extend([A,B])
                if s1 is s2:
                    b.is_circular_bead = True
                """
2480

2481
2482
2483
2484
        # """ Combine beads at junctions as needed """
        # for c,A,B in self.get_connections():
        #    ...

cmaffeo2's avatar
cmaffeo2 committed
2485
2486
2487
2488
2489
2490
2491
2492
        # ## Debug
        # all_beads = [b for s in segments for b in s.beads]
        # positions = np.array([b.position for b in all_beads])
        # dists = positions[:,np.newaxis,:] - positions[np.newaxis,:,:] 
        # ids = np.where( np.sum(dists**2,axis=-1) + 0.02**2*np.eye(len(dists)) < 0.02**2 )
        # print( ids )
        # pdb.set_trace()

2493
        """ Add intrahelical neighbors at connections """
2494
        special_seps = dict()
2495
2496
2497
        for c,A,B in self.get_connections("intrahelical"):
            b1,b2 = [l.particle for l in (A,B)]
            if b1 is b2:
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
                ## already handled by Segment._generate_beads, unless A,B are on same segment
                if A.container == B.container:
                    sorted_sites = sorted([(abs(L.address-b1.contour_position),L)
                                       for L in (A,B)])
                    close_site, far_site = [s[1] for s in sorted_sites]

                    b3 = far_site.container.get_nearest_bead(far_site.address)
                    b1.intrahelical_neighbors.append(b3)
                    b3.intrahelical_neighbors.append(b1)

                    if far_site.address > 0.5:
                        sep = b1.contour_position + (1-b3.contour_position)
                    else:
                        sep = b3.contour_position + (1-b1.contour_position)
                    sep = A.container.num_nt * sep
                    special_seps[(b1,b3)] = special_seps[(b3,b1)] = sep

2515
2516
            else:
                for b in (b1,b2): assert( b is not None )
cmaffeo2's avatar
cmaffeo2 committed
2517
                b1.make_intrahelical_neighbor(b2)
2518

2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
        def _combine_zero_sep_beads():
            skip_keys = []
            for b1,b2 in [k for k,v in special_seps.items() if v == 0]:
                if (b1,b2) in skip_keys: continue
                del special_seps[(b1,b2)]
                del special_seps[(b2,b1)]
                skip_keys.append((b2,b1))

                for b in b2.intrahelical_neighbors:
                    if b is b1: continue
                    if b is not b1 and b.parent is b2.parent:
                        sep = np.abs(b2.contour_position - b.contour_position) * b.parent.num_nt
                        special_seps[(b,b1)] = sep
                        special_seps[(b1,b)] = sep
                b1.combine(b2)

                for k in list(special_seps.keys()):
                    if b2 in k:
                        if b2 == k[0]:
                            newkey = (b1,k[1])
                        else:
                            newkey = (k[0],b2)
                        special_seps[newkey] = special_seps[k]
                        del special_seps[k]
        _combine_zero_sep_beads()

2545
        """ Reassign bead types """
2546
        if self.DEBUG: print("Assigning bead types")
2547
        beadtype_s = dict()
2548
2549
        beadtype_count = dict(D=0,O=0,S=0)

2550
2551
2552
        def _assign_bead_type(bead, num_nt, decimals):
            num_nt0 = bead.num_nt
            bead.num_nt = np.around( np.float32(num_nt), decimals=decimals )
2553
2554
            char = bead.type_.name[0].upper()
            key = (char, bead.num_nt)
2555
2556
2557
2558
            if key in beadtype_s:
                bead.type_ = beadtype_s[key]
            else:
                t = deepcopy(bead.type_)
2559
2560
2561
                t.__dict__["nts"] = bead.num_nt*2 if char in ("D","O") else bead.num_nt
                # t.name = t.name + "%03d" % (t.nts*10**decimals)
                t.name = char + "%03d" % (beadtype_count[char])
2562
                t.mass = t.nts * 150
2563
                t.diffusivity = 120 if t.nts == 0 else min( 50 / np.sqrt(t.nts/5), 120)
2564
                beadtype_count[char] += 1
cmaffeo2's avatar
cmaffeo2 committed
2565
                if self.DEBUG: print( "{} --> {} ({})".format(num_nt0, bead.num_nt, t.name) )
2566
2567
2568
2569
2570
                beadtype_s[key] = bead.type_ = t

        # (cluster_size[c-1])


2571
2572
2573
2574
        import scipy.cluster.hierarchy as hcluster
        beads = [b for s in segments for b in s if b.type_.name[0].upper() in ("D","O")]
        data = np.array([b.num_nt for b in beads])[:,np.newaxis]
        order = int(2-np.log10(2*max_basepairs_per_bead)//1)
2575
        try:
2576
            clusters = hcluster.fclusterdata(data, float(max_basepairs_per_bead)/500, criterion="distance")
2577
2578
            cluster_size = [np.mean(data[clusters == i]) for i in np.unique(clusters)]
        except:
2579
2580
            clusters = np.arange(len(data))+1
            cluster_size = data.flatten()
2581
        for b,c in zip(beads,clusters):
2582
            _assign_bead_type(b, cluster_size[c-1], decimals=order)
2583

2584
        beads = [b for s in segments for b in s if b.type_.name[0].upper() in ("S")]
2585
2586
        data = np.array([b.num_nt for b in beads])[:,np.newaxis]
        order = int(2-np.log10(max_nucleotides_per_bead)//1)
2587
        try:
2588
            clusters = hcluster.fclusterdata(data, float(max_nucleotides_per_bead)/500, criterion="distance")
2589
2590
            cluster_size = [np.mean(data[clusters == i]) for i in np.unique(clusters)]
        except:
2591
2592
            clusters = np.arange(len(data))+1
            cluster_size = data.flatten()
2593
        for b,c in zip(beads,clusters):
2594
            _assign_bead_type(b, cluster_size[c-1], decimals=order)
2595

2596
        self._apply_grid_potentials_to_beads(beadtype_s)
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
2608
2609
2610
2611
2612

        # for bead in [b for s in segments for b in s]:
        #     num_nt0 = bead.num_nt
        #     # bead.num_nt = np.around( np.float32(num_nt), decimals=decimals )
        #     key = (bead.type_.name[0].upper(), bead.num_nt)
        #     if key in beadtype_s:
        #         bead.type_ = beadtype_s[key]
        #     else:
        #         t = deepcopy(bead.type_)
        #         t.__dict__["nts"] = bead.num_nt*2 if t.name[0].upper() in ("D","O") else bead.num_nt
        #         # t.name = t.name + "%03d" % (t.nts*10**decimals)
        #         t.name = t.name + "%.16f" % (t.nts)
        #         print( "{} --> {} ({})".format(num_nt0, bead.num_nt, t.name) )
        #         beadtype_s[key] = bead.type_ = t


2613
2614
2615
        """ Update bead indices """
        self._countParticleTypes() # probably not needed here
        self._updateParticleOrder()
2616

2617
2618
2619
2620
2621
2622
2623

        def k_dsdna_bond(d):
            conversion = 0.014393265 # units "pN/AA" kcal_mol/AA^2
            elastic_modulus_times_area = 1000 # pN http://markolab.bmbcb.northwestern.edu/marko/Cocco.CRP.02.pdf
            return conversion*elastic_modulus_times_area/d

        def k_ssdna_bond(d):
2624
2625
2626
2627
            # conversion = 0.014393265 # units "pN/AA" kcal_mol/AA^2
            # ## TODO: get better numbers our ssDNA model
            # elastic_modulus_times_area = 800 # pN http://markolab.bmbcb.northwestern.edu/marko/Cocco.CRP.02.pdf
            
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
            return conversion*elastic_modulus_times_area/d

        def k_dsdna_angle(sep):
            Lp = get_effective_dsDNA_Lp(sep)
            return angle_spring_from_lp(sep,Lp)

        def k_xover_angle(sep):
            return 0.25 * angle_spring_from_lp(sep,147)


2638
        """ Add intrahelical bond potentials """
2639
        if self.DEBUG: print("Adding intrahelical bond potentials")
cmaffeo2's avatar
cmaffeo2 committed
2640
        dists = dict()          # intrahelical distances built for later use
2641
2642
2643
        intra_beads = self._get_intrahelical_beads() 
        if self.DEBUG: print("  Adding %d bonds" % len(intra_beads))
        for b1,b2 in intra_beads:
cmaffeo2's avatar
cmaffeo2 committed
2644
            # assert( not np.isclose( np.linalg.norm(b1.collapsedPosition() - b2.collapsedPosition()), 0 ) )
2645
            if np.linalg.norm(b1.collapsedPosition() - b2.collapsedPosition()) < 1:
cmaffeo2's avatar
cmaffeo2 committed
2646
2647
                # print("WARNING: some beads are very close")
                ...
2648

cmaffeo2's avatar
cmaffeo2 committed
2649
            parent = self._getParent(b1,b2)
2650

2651
2652
2653
2654
2655
2656
            if (b1,b2) in special_seps:
                sep = special_seps[(b1,b2)]
            else:
                seg = b2.parent
                c0 = b2.contour_position
                sep = np.abs(b1.get_nt_position(seg,c0)-b2.get_nt_position(seg))
cmaffeo2's avatar
cmaffeo2 committed
2657

2658
2659
            is_dsdna = b1.type_.name[0] == "D" and b2.type_.name[0] == "D"
            if is_dsdna:
2660
2661
                d = 3.4*sep
            else:
2662
2663
                # d = 6.5*sep     # https://journals.aps.org/pre/abstract/10.1103/PhysRevE.86.021901
                d = 6.4*sep     # https://journals.aps.org/pre/abstract/10.1103/PhysRevE.86.021901
2664
2665
2666
2667
                if b1.type_.name[0] != b2.type_.name[0]:
                    """ Add a small extra distance to junction """
                    d += 3

2668
            if b1 not in dists:
cmaffeo2's avatar
cmaffeo2 committed
2669
                dists[b1] = dict()
2670
            if b2 not in dists:
cmaffeo2's avatar
cmaffeo2 committed
2671
2672
2673
2674
2675
                dists[b2] = dict()
            # dists[b1].append([b2,sep])
            # dists[b2].append([b1,sep])
            dists[b1][b2] = sep
            dists[b2][b1] = sep
2676

2677
2678
            if b1 is b2: continue

2679
            # dists[[b1,b2]] = dists[[b2,b1]] = sep
2680
2681
            if is_dsdna:
                k = k_dsdna_bond(d)
2682
                bond = self.get_bond_potential(k,d,correct_geometry=True)
2683
2684
            else:
                bond = self._get_wlc_sk_bond_potential(d)
2685
            parent.add_bond( b1, b2, bond, exclude=True )
2686

2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
2698
2699
2700
        # for s in self.segments:
        #     sepsum = 0
        #     beadsum = 0
        #     for b1 in s.beads:
        #         beadsum += b1.num_nt
        #         for bead_list in self._recursively_get_beads_within_bonds(b1, 1):
        #             assert(len(bead_list) == 1)
        #             if b1.idx < bead_list[-1].idx: # avoid double-counting
        #                 for b2 in bead_list:
        #                     if b2.parent == b1.parent:
        #                         sepsum += dists[b1][b2]
        #     sepsum += sep
        #     print("Helix {}: bps {}, beads {}, separation {}".format(s.name, s.num_nt, beadsum, sepsum))

2701
        """ Add intrahelical angle potentials """
2702
2703
2704
2705
2706
2707
2708
2709
2710
2711
2712
2713
2714
2715
        def get_effective_dsDNA_Lp(sep):

            """ The persistence length of our model was found to be a
            little off (probably due to NB interactions). This
            attempts to compensate """

            ## For 1 bp, Lp=559, for 25 Lp = 524
            beads_per_bp = sep/2
            Lp0 = 147
            # return 0.93457944*Lp0 ;# factor1
            return 0.97*Lp0 ;# factor2
            # factor = bead_per_bp * (0.954-0.8944
            # return Lp0 * bead_per_bp

2716
        if self.DEBUG: print("Adding intrahelical angle potentials")
2717
        for b1,b2,b3 in self._get_intrahelical_angle_beads():
cmaffeo2's avatar
cmaffeo2 committed
2718
            parent = self._getParent(b1,b2,b3)
2719
2720
            seg = b2.parent
            c0 = b2.contour_position
2721
            sep = dists[b2][b1] + dists[b2][b3]
2722

2723

2724
            if b1.type_.name[0] == "D" and b2.type_.name[0] == "D" and b3.type_.name[0] == "D":
2725
2726
                
                k = k_dsdna_angle(sep)
cmaffeo2's avatar
cmaffeo2 committed
2727
                if local_twist:
2728
2729
2730
2731
                    k_dihed = 0.25*k
                    k *= 0.75    # reduce because orientation beads impose similar springs
                    dihed = self.get_dihedral_potential(k_dihed,180)
                    parent.add_dihedral(b1,b2,b2.orientation_bead,b3, dihed)
2732
                angle = self.get_angle_potential(k,180)
2733

2734

2735
            elif b1.type_.name[0] == "S" and b2.type_.name[0] == "S" and b3.type_.name[0] == "S":
2736
2737
                # angle = self._get_wlc_sk_angle_potential(sep*6.5) # TODO move 6.5 somewhere sensible
                angle = self._get_wlc_sk_angle_potential(sep*6.4) # TODO move 6.4 somewhere sensible
2738
2739
2740
            else:
                ## Considered as a sscrossover below
                continue
2741

2742
2743
2744
            parent.add_angle( b1, b2, b3, angle )

        """ Add intrahelical exclusions """
2745
        if self.DEBUG: print("Adding intrahelical exclusions")
2746
        beads = dists.keys()
cmaffeo2's avatar
cmaffeo2 committed
2747
        def _recursively_get_beads_within(b1,d,done=()):
2748
            ret = []
cmaffeo2's avatar
cmaffeo2 committed
2749
            for b2,sep in dists[b1].items():
2750
2751
2752
2753
2754
2755
2756
                if b2 in done: continue
                if sep < d:
                    ret.append( b2 )
                    done.append( b2 )
                    tmp = _recursively_get_beads_within(b2, d-sep, done)
                    if len(tmp) > 0: ret.extend(tmp)
            return ret
2757

2758
2759
        exclusions = set()
        for b1 in beads:
2760
2761
2762
2763
2764
2765
2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
            """ In addition to bond exclusiosn, only add exclusions
            within like-type segments (i.e. dsDNA or ssDNA, not
            junctions between the two) """

            t = type(b1.parent)
            if t is DoubleStrandedSegment:
                cutoff = 20
            elif t is SingleStrandedSegment:
                cutoff = 5
            else:
                raise ValueError("Unexpected polymer segment type")

            for b in _recursively_get_beads_within(b1, cutoff, done=[b1]):
                if isinstance(b.parent,t):
                    exclusions.add((b1,b))
                else:
                    break
            # exclusions.update( tmp )
2778

2779
2780
2781
2782
2783
2784
2785
2786
2787
        ## Add exclusions very near crossovers
        for c,A,B in sum([self.get_connections(term) for term in ("crossover","terminal_crossover")],[]):
            b1,b2 = [loc.particle for loc in (A,B)]
            near_b1 = _recursively_get_beads_within(b1, 3, done=[])
            near_b2 = _recursively_get_beads_within(b2, 3, done=[])
            for bi in near_b1:
                for bj in near_b2:
                    exclusions.add((bi,bj))

2788
        if self.DEBUG: print("Adding %d exclusions" % len(exclusions))
2789
        for b1,b2 in exclusions:
cmaffeo2's avatar
cmaffeo2 committed
2790
            parent = self._getParent(b1,b2)
2791
            parent.add_exclusion( b1, b2 )
cmaffeo2's avatar
cmaffeo2 committed
2792
2793
2794

        """ Twist potentials """
        if local_twist:
2795
            if self.DEBUG: print("Adding twist potentials")
cmaffeo2's avatar
cmaffeo2 committed
2796
2797
2798

            for b1 in beads:
                if "orientation_bead" not in b1.__dict__: continue
cmaffeo2's avatar
cmaffeo2 committed
2799
                for b2,sep in dists[b1].items():
cmaffeo2's avatar
cmaffeo2 committed
2800
                    if "orientation_bead" not in b2.__dict__: continue
2801
                    if b2.idx < b1.idx: continue # Don't double-count
cmaffeo2's avatar
cmaffeo2 committed
2802
2803
2804
2805
2806
2807
2808

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

                    parent = self._getParent( b1, b2 )

                    """ Add heuristic 90 degree potential to keep orientation bead orthogonal """
2809
                    k = 0.25*k_dsdna_angle(sep)
cmaffeo2's avatar
cmaffeo2 committed
2810
2811
2812
                    pot = self.get_angle_potential(k,90)
                    parent.add_angle(o1,b1,b2, pot)
                    parent.add_angle(b1,b2,o2, pot)
2813

cmaffeo2's avatar
cmaffeo2 committed
2814
2815
2816
2817
                    ## TODO: improve this
                    twist_per_nt = 0.5 * (p1.twist_per_nt + p2.twist_per_nt)
                    angle = sep*twist_per_nt
                    if angle > 360 or angle < -360:
cmaffeo2's avatar
cmaffeo2 committed
2818
2819
                        print("WARNING: twist angle out of normal range... proceeding anyway")
                        # raise Exception("The twist between beads is too large")
cmaffeo2's avatar
cmaffeo2 committed
2820
2821
                        
                    k = self._get_twist_spring_constant(sep)
2822
2823
2824
2825
                    if escapable_twist:
                        pot = self.get_dihedral_potential(k,angle,max_potential=1)
                    else:
                        pot = self.get_dihedral_potential(k,angle)
cmaffeo2's avatar
cmaffeo2 committed
2826
2827
                    parent.add_dihedral(o1,b1,b2,o2, pot)

2828

2829
2830
2831
2832
2833
        def count_crossovers(beads):
            count = 0
            for b in beads:
                for l in b.locations:
                    if l.connection is not None:
2834
                        if l.connection.type_ in ("crossover","terminal_crossover"):
2835
2836
2837
                            count += 1
            return count

2838
        def add_local_crossover_strand_orientation_potential(b1,b2, b1_on_fwd_strand):
2839

2840
2841
            """ Adds a dihedral angle potential so bead b2 at opposite
            end of crossover stays on correct side of helix of b1 """
2842

2843
2844
            u1 = b1.get_intrahelical_above(all_types=False)
            d1 = b1.get_intrahelical_below(all_types=False)
2845

2846
            sign = 1 if b1_on_fwd_strand else -1
2847

cmaffeo2's avatar
cmaffeo2 committed
2848
            # if b1.parent.name == "8-1" or b2.parent.name == "8-1":
2849
2850
2851
2852
            #     print()
            #     print(b1.parent.name, b2.parent.name, b1_on_fwd_strand)
            #     import pdb
            #     pdb.set_trace()
2853

2854
2855
2856
2857
2858
2859
2860
2861
2862
            a,b,c = b2,b1,d1
            if c is None or c is a:
                c = u1
                sign *= -1
            if c is None or c is a: return
            try:
                d = b1.orientation_bead
            except:
                return
2863

2864
2865
2866
            k = k_xover_angle(sep=1) # TODO
            pot = self.get_dihedral_potential(k, sign*120)
            self.add_dihedral( a,b,c,d, pot )
2867

2868
        def add_local_tee_orientation_potential(b1,b2, b1_on_fwd_strand, b2_on_fwd_strand):
2869

2870
2871
2872
            """ b1 is the end of a helix, b2 is in the middle This
            adds a dihedral angle potential so helix of b1 is oriented
            properly relative to strand on b2 """
2873

2874
2875
            u1,u2 = [b.get_intrahelical_above(all_types=False) for b in (b1,b2)]
            d1,d2 = [b.get_intrahelical_below(all_types=False) for b in (b1,b2)]
2876

2877
2878
            angle = 150
            if not b2_on_fwd_strand: angle -= 180
2879

2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
            a,b,c = u2,b2,b1
            if a is None:
                a = d2
                angle -= 180
            try:
                d = b1.orientation_bead
            except:
                d = None
            angle -= 120

            while angle > 180:
                angle -= 360
            while angle < -180:
                angle += 360

            k = k_xover_angle(sep=1) # TODO
            if a is not None and d is not None:
                pot = self.get_dihedral_potential(k,angle)
                self.add_dihedral( a,b,c,d, pot )

            ## Add 180 degree angle potential
            a,b,c = b2,b1,u1
            if c is None: c = d1

            if c is not None:
                pot = self.get_angle_potential(0.5*k,180)
                self.add_angle( a,b,c, pot )

        def add_parallel_crossover_potential(b1,b2):
cmaffeo2's avatar
cmaffeo2 committed
2909

2910
            ## Get beads above and below
2911
2912
            u1,u2 = [b.get_intrahelical_above(all_types=False) for b in (b1,b2)]
            d1,d2 = [b.get_intrahelical_below(all_types=False) for b in (b1,b2)]
2913
2914
2915
2916
2917
2918
            dotProduct = b1.parent.contour_to_tangent(b1.contour_position).dot(
                b2.parent.contour_to_tangent(b2.contour_position) )
            if dotProduct < 0:
                tmp = d2
                d2  = u2
                u2  = tmp
2919

2920
            a = None
2921
            if u1 is not None and u2 is not None:
2922
                t0 = self.hj_equilibrium_angle
2923
                a,b,c,d = (u1,b1,b2,u2)
2924
            elif d1 is not None and d2 is not None:
2925
                t0 = self.hj_equilibrium_angle
2926
                a,b,c,d = (d1,b1,b2,d2 )
2927
            elif d1 is not None and u2 is not None:
2928
                t0 = self.hj_equilibrium_angle-180
2929
                a,b,c,d = (d1,b1,b2,u2)
2930
            elif u1 is not None and d2 is not None:
2931
                t0 = self.hj_equilibrium_angle-180
2932
2933
                a,b,c,d = (u1,b1,b2,d2)

2934
2935
2936
2937
2938
            if t0 > 180:
                t0 = t0-360
            elif t0 < -180:
                t0 = t0+360

2939
            ## TODO?: Check length-dependence of this potential
2940
            if a is not None:
2941
2942
2943
2944
2945
2946
2947
                k_intrinsic = 0.00086
                k = [1/k for k in (4*k_xover_angle( dists[b][a] ),
                                  k_intrinsic,
                                  4*k_xover_angle( dists[c][d] ))]
                k = sum(k)**-1
                k = k * count_crossovers((b1,b2))/4

2948
                pot = self.get_dihedral_potential(k,t0)
2949
                self.add_dihedral( a,b,c,d,  pot )
2950
2951
2952
2953
            ...

        """ Functions for adding crossover potentials  """
        def add_ss_crossover_potentials(connection,A,B, add_bond=True):
cmaffeo2's avatar
cmaffeo2 committed
2954
            b1,b2 = [loc.particle for loc in (A,B)]
2955
2956
2957
2958
2959
2960
2961
2962
2963
2964
2965
2966
2967
2968
2969

            if (b1,b2,A.on_fwd_strand,B.on_fwd_strand) in processed_crossovers:
                return
            processed_crossovers.add((b1,b2,A.on_fwd_strand,B.on_fwd_strand))
            processed_crossovers.add((b2,b1,B.on_fwd_strand,A.on_fwd_strand))

            if b1 is b2:
                """ Catch attempts to add "crossover potentials" at
                intrahelical junctions between ds and ssDNA """
                if A.container is not b1.parent:
                    b1 = A.container.get_nearest_bead(A.address)
                if B.container is not b2.parent:
                    b2 = B.container.get_nearest_bead(B.address)
                if b1 is b2:
                    return
2970
2971
2972

            if b1 is None or b2 is None:
                return
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986

            ## TODO: improve parameters
            if add_bond:
                pot = self.get_bond_potential(4,12)
                self.add_bond(b1,b2, pot)

            ## Add potentials to provide a sensible orientation
            ## TODO refine potentials against all-atom simulation data
            if local_twist:
                add_local_crossover_strand_orientation_potential(b1,b2, A.on_fwd_strand)
                add_local_crossover_strand_orientation_potential(b2,b1, B.on_fwd_strand)

        def add_crossover_potentials(connection,A,B):
            ## TODO: use a better description here
cmaffeo2's avatar
cmaffeo2 committed
2987
            b1,b2 = [loc.particle for loc in (A,B)]
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999
3000
            if (b1,b2,A.on_fwd_strand,B.on_fwd_strand) in processed_crossovers:
                return

            processed_crossovers.add((b1,b2,A.on_fwd_strand,B.on_fwd_strand))
            processed_crossovers.add((b2,b1,B.on_fwd_strand,A.on_fwd_strand))

            if b1 is b2:
                """ Catch attempts to add "crossover potentials" at
                intrahelical junctions between ds and ssDNA """
                return

            """ Add parallel helices potential, possibly """
            ## Add potential to provide a particular orinetation
For faster browsing, not all history is shown. View entire blame