polygon_mesh.py 18.9 KB
Newer Older
cmaffeo2's avatar
cmaffeo2 committed
1
2
import numpy as np
import sys
3
import re
cmaffeo2's avatar
cmaffeo2 committed
4
from coords import rotationAboutAxis
5

6
7
from segmentmodel import SegmentModel, SingleStrandedSegment, DoubleStrandedSegment

8
9
10
11
12
13
14
15
16
class MayaObj():
    def __init__(self, maya_lines):
        self.name, self.parent_name = MayaObj._parse_first_maya_line( maya_lines[0] )

    def _parse_first_maya_line(l):
        ## createNode vHelix -n "helix_1";
        name = parent = None
        vals = l.split()
        for i in range(len(vals)):
17
18
19
20
            if vals[i] == "-n": 
                name = MayaObj._strip_maya_string( vals[i+1] )
            if vals[i] == "-p":
                parent = MayaObj._strip_maya_string( vals[i+1] )
21
        return name,parent
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
    
    def _strip_maya_string(string):
        return re.match('"([a-zA-Z_[0-9]+)";?$', string).group(1)


class MayaConnection():
    def __init__(self, helix1, base1, suff1, helix2, base2, suff2):
        self.helix1 = helix1
        self.base1  = base1
        self.suff1  = suff1
        self.helix2 = helix2
        self.base2  = base2
        self.suff2  = suff2

# MayaConnections = dict()        

def ParseMayaConnection(line, base_dict):
    words = line.split()
    if len(words) != 3:
        return None
    cmd,obj1,obj2 = words

    def parse_obj(obj):
        ## Typically, base name is not unique, so parent name used too
        result = re.match('"(.*(forw)?(backw)?_[0-9]+)\.([a-z]{2})"', obj)
        if result is not None:
            name, tmp, tmp, suff = result.groups()
            if name in base_dict:
                return base_dict[name].parent.name, name, suff
            else:
                result = re.match('\|(.*)\|(.*)', name)
                if result is not None:
                    helix, name = result.groups()
                    return helix, name, suff

        return None, None, None
                    
    helix1,base1,suff1 = parse_obj(obj1)
    if helix1 is not None:
        helix2,base2,suff2 = parse_obj(obj2)
        if helix2 is not None:
            return MayaConnection(helix1,base1,suff1,helix2,base2,suff2)
    return None

66
67
68
69

class MayaHelix(MayaObj):
    def __init__(self, maya_lines):
        MayaObj.__init__(self, maya_lines)
70
71
        # self.bases = []
        self.bases = dict()
cmaffeo2's avatar
cmaffeo2 committed
72
73
74
75
76
77
        self.position = None
        self.orientation = None

        self._parse_maya_lines(maya_lines)
        
    def _parse_maya_lines(self, maya_lines):
78
        """
cmaffeo2's avatar
cmaffeo2 committed
79
80
81
82
createNode vHelix -n "helix_30";
setAttr ".t" -type "double3" 9.24809 21.4638 -12.465 ;
setAttr ".r" -type "double3" 28.99971393816978 2.2177410772603801 87.637950315471372 ;
setAttr ".dh" yes;
83
        """
cmaffeo2's avatar
cmaffeo2 committed
84
85
86
87
        for l in maya_lines:
            vals = l.split()
            if vals[0] == "setAttr":
                if vals[1] == '".t"':
88
                    self.position = 10*np.array([float(a) for a in vals[4:7]])
cmaffeo2's avatar
cmaffeo2 committed
89
90
91
92
93
94
                elif vals[1] == '".r"':
                    # self.eulerAngles = np.array([float(a) for a in vals[4:7]]) # defined here
                    ax,ay,az = [float(a) for a in vals[4:7]] 
                    Rx = rotationAboutAxis( [1,0,0], ax, normalizeAxis=False )
                    Ry = rotationAboutAxis( [0,1,0], ay, normalizeAxis=False )
                    Rz = rotationAboutAxis( [0,0,1], az, normalizeAxis=False )
95
96
                    self.orientation = Rz.dot(Ry.dot(Rx)) # TODO: check
                    # rotXYZ = np.array()
cmaffeo2's avatar
cmaffeo2 committed
97

98
99
100
101
    def add_base(self, base):
        # self.bases.append(base)
        self.bases[base.name] = base
        base.parent = self
102

103
104
105
106
107
108
109
110
111
112
113
114
115
116
    def find_double_helix_regions(self):
        double_helices = []
        for bn,b1 in self.bases.items():
            b2 = b1.basepair
            if b2 is None: continue
            d =  np.linalg.norm( b1.get_position()-b2.get_position() )
            print( d )
            

    def write_xyz(self,file_handle = sys.stdout):
        for bn,b in self.bases.items():
            position = b.get_position()
            # position = b.position.dot( self.orientation ) + self.position # TODO: double check
            file_handle.write("%s %f %f %f\n" % (b.name[0], position[0], position[1], position[2]))
117
118


cmaffeo2's avatar
cmaffeo2 committed
119
class MayaBase(MayaObj):
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
    def __init__(self, maya_lines):
        MayaObj.__init__(self, maya_lines)
        self._parse_maya_lines(maya_lines)        
        self.basepair = None
        self.parent = None
        self.end3 = None
        self.end5 = None

    def _parse_maya_lines(self, maya_lines):
        """
createNode HelixBase -n "backw_1" -p "helix_30";
setAttr ".t" -type "double3" 0.36767788771440857 -0.78848777472188547 -6.014 ;
        """
        for l in maya_lines:
            vals = l.split()
            if vals[0] == "setAttr":
                if vals[1] == '".t"':
                    self.position = 10*np.array([float(a) for a in vals[4:7]])
138

139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
    def add_basepair(self, base):
        assert( self.parent == base.parent )
        assert( self.basepair is None or self.basepair == base )
        assert( base.basepair is None or base.basepair == self )
        self.basepair = base
        base.basepair = self

    def add_end3(self, base):
        assert( self.end3 is None or self.end3 == base )
        assert( base.end5 is None or base.end5 == self )
        self.end3 = base
        base.end5 = self

    def add_end5(self, base):
        base.add_end3(self)
    
    def get_position(self):
        h = self.parent
        return h.orientation.dot(self.position) + h.position # TODO: double check
cmaffeo2's avatar
cmaffeo2 committed
158
        
159

160
161
162
163
# class MayaAimConstraint(MayaObj):
#     ...


cmaffeo2's avatar
cmaffeo2 committed
164
# set ch [open /home/cmaffeo2/Downloads/ball.ma]
165
166


cmaffeo2's avatar
cmaffeo2 committed
167
168
169
170
171
def parse_maya_file(maya_file):

    helices = dict()
    bases = []
    aimConstraints = []
172
    connections = []
cmaffeo2's avatar
cmaffeo2 committed
173
174
175
176
177

    ## Parse through .ma file
    with open(maya_file) as fh:
        linesBuffer = []
        for line in fh:
178
179
180
            words = line.split()
            if words[0] == "createNode":

cmaffeo2's avatar
cmaffeo2 committed
181
182
183
184
                ## Create a new object
                if len(linesBuffer) > 0:
                    objType = linesBuffer[0].split()[1]
                    if objType == "vHelix":
185
                        h = MayaHelix( linesBuffer )
cmaffeo2's avatar
cmaffeo2 committed
186
187
188
                        helices[h.name] = h
                    elif objType == "HelixBase":
                        bases.append( MayaBase( linesBuffer ) )
189
190
                    # elif objType == "aimConstraint":
                    #     aimConstraints.append( MayaAimConstraint( linesBuffer ) )
cmaffeo2's avatar
cmaffeo2 committed
191
192
                ## Clear lines buffer
                linesBuffer = []
193
194
195
            elif words[0] == "connectAttr":
                connections.append(line)
                
cmaffeo2's avatar
cmaffeo2 committed
196
            ## Extend lines buffer
197
198
            linesBuffer.append(line)
            
cmaffeo2's avatar
cmaffeo2 committed
199
200
    ## Crawl through bases, attach to parent helix
    for b in bases:
201
        helix_name = b.parent_name
cmaffeo2's avatar
cmaffeo2 committed
202
        helices[helix_name].add_base( b ) # function updates base position?
203
204
205
206
207
208
209
210
211

    ## Parse connections
    new_connections = []
    base_dict = { b.name:b for b in bases }
    for line in connections:        
        conn = ParseMayaConnection(line, base_dict)
        if conn is not None:
            new_connections.append(conn)
    connections = new_connections
cmaffeo2's avatar
cmaffeo2 committed
212
        
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
    ## Assign base connectivity
    for c in connections:
        h1 = c.helix1
        bn1 = c.base1
        s1 = c.suff1
        h2 = c.helix2
        bn2 = c.base2
        s2 = c.suff2
        
        b1 = helices[h1].bases[bn1]
        b2 = helices[h2].bases[bn2]
        if s1 == "lb" and s2 == "lb":
            b1.add_basepair(b2)
        # elif (bn1[0] == "f" and s1 == "fw" and s2 == "bw") or \
        #      (bn1[0] == "b" and s1 == "bw" and s2 == "fw"):
        #     b1.add_end3(b2)
        # elif (bn1[0] == "f" and s1 == "bw" and s2 == "fw") or \
        #      (bn1[0] == "b" and s1 == "fw" and s2 == "bw"):
        #     b1.add_end5(b2)
        elif s1 == "fw" and s2 == "bw":
            b1.add_end3(b2)
        elif s1 == "bw" and s2 == "fw":
            b1.add_end5(b2)
        else:
            raise Exception("Unkown connection %s %s" % (s1,s2))
cmaffeo2's avatar
cmaffeo2 committed
238

239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
    # # base_positions = np.zeros( [ len(bases), 3 ] )
    # base_positions = np.array( [b.position for b in bases] )
    # for aim in aimConstraints:
    #     base = base_dict[ aim.parent_name ]

    return helices

def convert_maya_to_segments(helices):

    def get_end3_ssDNA(nt):
        ret = []
        if nt.end3 is not None:
            # print( nt.parent.name, nt.end3.parent.name, nt.end3.name )
            nt2 = nt.end3
            if nt2.basepair is None:
                ret.append(nt2)
                ret.extend(get_end3_ssDNA(nt2))
        return ret

    def get_end5_ssDNA(nt):
        ret = []
        if nt.end5 is not None:
            nt2 = nt.end5
            if nt2.basepair is None:
                ret.append(nt2)
                ret.extend(get_end5_ssDNA(nt2))
        return ret
cmaffeo2's avatar
cmaffeo2 committed
266

267
268
269
270
271
272
273
274
275
    """ Build dsDNA helices """
    segments = dict()
    segmentsBot = dict()
    segmentsTop = dict()
    for hn,h in helices.items():
        paired_bases = [b for bn,b in h.bases.items() if b.basepair is not None]
        upaired = [b for bn,b in h.bases.items() if b.basepair is None]
        # for b in upaired:
        #     print(hn, b.name, b.end3, b.end5)
cmaffeo2's avatar
cmaffeo2 committed
276
        
277
278
279
        z = [b.position[2] for b in paired_bases]
        bot = paired_bases[ np.argmin(z) ]
        top = paired_bases[ np.argmax(z) ]
cmaffeo2's avatar
cmaffeo2 committed
280

281
282
283
284
285
286
        if bot.name[0] == "f": bot = bot.basepair
        if top.name[0] == "b": top = top.basepair

        botPos = 0.5 * (bot.get_position() + bot.basepair.get_position())
        topPos = 0.5 * (top.get_position() + top.basepair.get_position())

287
        segments[hn] =  DoubleStrandedSegment(hn, num_bp = len(z)//2,
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
                                              start_position = botPos,
                                              end_position = topPos)
        segmentsBot[hn] = bot
        segmentsTop[hn] = top

    """ Sanity check """
    def end_crosses(end,direction):
        assert(direction in ("end3","end5"))
        other = end.__dict__[direction]
        if other is None or end.parent != other.parent:
            return 1
        else:
            return 0

    hs = segments.keys()
    botCrossings1 = [end_crosses(segmentsBot[h],"end3") for h in hs]
    botCrossings2 = [end_crosses(segmentsBot[h].basepair,"end5") for h in hs]
    topCrossings1 = [end_crosses(segmentsTop[h],"end3") for h in hs]
    topCrossings2 = [end_crosses(segmentsTop[h].basepair,"end5") for h in hs]
    # print( [np.mean(a) for a in (botCrossings1,botCrossings2,topCrossings1,topCrossings2)] )

    # for a in (botCrossings1,botCrossings2,topCrossings1,topCrossings2:
    #     print(np.mean(a))
    #     assert( np.mean(a) > 0.5 )

    """ Build ssDNA segments """
    ss_segments = []
    connectionToSingleStrand = dict()
    for hn,seg in segments.items():
        bot = segmentsBot[hn]
        top = segmentsTop[hn]
        # if top.end3 is not None:
        #     print(top.parent.name,top.name, top.end3.parent.name, top.end3.name)
        # if bot.end3 is not None:
        #     print(bot.parent.name,bot.name, bot.end3.parent.name, bot.end3.name)

        ## TODO: clean this up

        if top not in connectionToSingleStrand:
            ssSeg = get_end3_ssDNA(top)
            if len(ssSeg) > 0:
                seg2 = SingleStrandedSegment("strand", 
                                              start_position = ssSeg[0].get_position(),
                                              end_position = ssSeg[-1].get_position(),
                                              num_nts = len(ssSeg))
                seg.connect_end3(seg2)
                connectionToSingleStrand[top] = seg2
                other = ssSeg[-1].end3
                if other is not None:
                    hn3 = other.parent.name
                    seg3 = segments[hn3]
                    bot3 = segmentsBot[hn3]
                    top3 = segmentsTop[hn3]
                    # if bot3 == other or bot3.basepair == other:
                    #     seg3.connect_start5(seg2)
                    # elif top3 == other or top3.basepair == other:
                    #     seg3.connect_end5(seg2)
                    if bot3.basepair == other:
                        seg3.connect_start5(seg2)
                    elif top3.basepair == other:
                        seg3.connect_end5(seg2)
                    else:
                        raise Exception("BUG")
                    connectionToSingleStrand[other] = seg2
                ss_segments.append(seg2)

        if top.basepair not in connectionToSingleStrand:
            ssSeg = get_end5_ssDNA(top.basepair)[::-1]
            if len(ssSeg) > 0:
                seg2 = SingleStrandedSegment("strand", 
                                              start_position = ssSeg[-1].get_position(),
                                              end_position = ssSeg[0].get_position(),
                                              num_nts = len(ssSeg))
                seg.connect_end5(seg2)
                connectionToSingleStrand[top.basepair] = seg2
                other = ssSeg[0].end5
                if other is not None:
                    hn3 = other.parent.name
                    seg3 = segments[hn3]
                    bot3 = segmentsBot[hn3]
                    top3 = segmentsTop[hn3]
                    # if bot3 == other or bot3.basepair == other:
                    #     seg3.connect_start5(seg2)
                    # elif top3 == other or top3.basepair == other:
                    #     seg3.connect_end5(seg2)
                    if bot3 == other:
                        seg3.connect_start3(seg2)
                    elif top3 == other:
                        seg3.connect_end3(seg2)
                    else:
                        raise Exception("BUG")
                    connectionToSingleStrand[other] = seg2
                ss_segments.append(seg2)

        if bot not in connectionToSingleStrand:
            ssSeg = get_end3_ssDNA(bot)
            if len(ssSeg) > 0:
                seg2 = SingleStrandedSegment("strand", 
                                              start_position = ssSeg[0].get_position(),
                                              end_position = ssSeg[-1].get_position(),
                                              num_nts = len(ssSeg))
                seg.connect_start3(seg2)
                connectionToSingleStrand[bot] = seg2
                other = ssSeg[-1].end3
                if other is not None:
                    hn3 = other.parent.name
                    seg3 = segments[hn3]
                    bot3 = segmentsBot[hn3]
                    top3 = segmentsTop[hn3]
                    if bot3.basepair == other:
                        seg3.connect_start5(seg2)
                    elif top3.basepair == other:
                        seg3.connect_end5(seg2)
                    else:
                        raise Exception("BUG")
                    connectionToSingleStrand[other] = seg2
                ss_segments.append(seg2)

        if bot.basepair not in connectionToSingleStrand:
            ssSeg = get_end5_ssDNA(bot.basepair)[::-1]
            if len(ssSeg) > 0:
                seg2 = SingleStrandedSegment("strand", 
                                              start_position = ssSeg[-1].get_position(),
                                              end_position = ssSeg[0].get_position(),
                                              num_nts = len(ssSeg))
                seg.connect_start5(seg2)
                connectionToSingleStrand[bot.basepair] = seg2
                other = ssSeg[0].end5
                if other is not None:
                    hn3 = other.parent.name
                    seg3 = segments[hn3]
                    bot3 = segmentsBot[hn3]
                    top3 = segmentsTop[hn3]
                    if bot3 == other:
                        seg3.connect_start3(seg2)
                    elif top3 == other:
                        seg3.connect_end3(seg2)
                    else:
                        raise Exception("BUG")
                    connectionToSingleStrand[other] = seg2
                ss_segments.append(seg2)
                
    ## Connect remaining dsDNA ends
    no_connection_count = 0
    for hn,seg in segments.items():
        bot = segmentsBot[hn]
        top = segmentsTop[hn]

        top_bp, bot_bp = [b.basepair for b in (top,bot)]
        
        if top not in connectionToSingleStrand and top.end3 is not None:
            b2 = top.end3
            h2 = b2.parent.name
            s2 = segments[h2]
            if b2 == segmentsBot[h2] or b2.basepair == segmentsBot[h2]:
                seg.connect_end3(s2.start5,type_="terminal_crossover")
            elif b2 == segmentsTop[h2] or b2.basepair == segmentsTop[h2]:
                seg.connect_end3(s2.end5,type_="terminal_crossover")
            else:
                raise Exception("BUG")
        
        if top_bp not in connectionToSingleStrand and top_bp.end5 is not None:
            b2 = top_bp.end5
            h2 = b2.parent.name
            s2 = segments[h2]
            if b2 == segmentsBot[h2] or b2.basepair == segmentsBot[h2]:
                seg.connect_end5(s2.start3,type_="terminal_crossover")
            elif b2 == segmentsTop[h2] or b2.basepair == segmentsTop[h2]:
                seg.connect_end5(s2.end3, type_="terminal_crossover")
            else:
                raise Exception("BUG")

        if bot not in connectionToSingleStrand and bot.end3 is not None:
            b2 = bot.end3
            h2 = b2.parent.name
            s2 = segments[h2]
            if b2 == segmentsBot[h2] or b2.basepair == segmentsBot[h2]:
                seg.connect_start3(s2.start5,type_="terminal_crossover")
            elif b2 == segmentsTop[h2] or b2.basepair == segmentsTop[h2]:
                seg.connect_start3(s2.end5,type_="terminal_crossover")
            else:
                raise Exception("BUG")
        
        if bot_bp not in connectionToSingleStrand and bot_bp.end5 is not None:
            b2 = bot_bp.end5
            h2 = b2.parent.name
            s2 = segments[h2]
            if b2 == segmentsBot[h2] or b2.basepair == segmentsBot[h2]:
                seg.connect_start5(s2.start3,type_="terminal_crossover")
            elif b2 == segmentsTop[h2] or b2.basepair == segmentsTop[h2]:
                seg.connect_start5(s2.end3, type_="terminal_crossover")
            else:
                raise Exception("BUG")

    return [s for n,s in segments.items()] + ss_segments, segments


def write_xyz(filename, helices):
    nbeads = 0
    for hn,h in helices.items():
        nbeads += len(h.bases)

    with open("test.xyz","w") as fh:
        fh.write("%d\n\n" % nbeads)
        for hn,h in helices.items():
            h.write_xyz(fh)
            #      h.find_double_stranded_regions


if __name__ == "__main__":
    maya_helices = parse_maya_file("/home/cmaffeo2/Downloads/bunny.ma")

    segments,dsSegmentDict = convert_maya_to_segments(maya_helices)
    
    model = SegmentModel( segments,
                          local_twist = True,
                          max_basepairs_per_bead = 1,
                          max_nucleotides_per_bead = 1,
                          dimensions=(5000,5000,5000), DEBUG=True )

    model.simulate( outputPrefix = 'bunny2', outputPeriod=1e4, numSteps=1e8, gpu=1 )
cmaffeo2's avatar
cmaffeo2 committed
509