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

6
from ..version import maintainer
7
from ..segmentmodel import SegmentModel, SingleStrandedSegment, DoubleStrandedSegment
8

9
10
11
12
DEBUG=False

maya_obj_registry = dict()
maya_obj_registry_short = dict()
13
class MayaObj():
14
15
    """ Class representing a node in a Maya ascii file """

cmaffeo2's avatar
cmaffeo2 committed
16
17
    seq_mapping = {str(i):c for c,i in zip('ATCG',range(4))}

18
    def __init__(self, maya_lines):
19
20
21
22
23
        self.parent = None
        self.children = dict()
        self.position = np.array((0,0,0))
        self.orientation = np.eye(3)

24
        self.name, self.parent_name = MayaObj._parse_first_maya_line( maya_lines[0] )
25
26
27
28
29
30
31
32
33
34
35
36
37
38
        self._parse_maya_lines(maya_lines)

        if self.parent_name is not None:
            if self.parent_name in maya_obj_registry:
                p = maya_obj_registry[self.parent_name]
                p.add(self)
            elif self.parent_name in maya_obj_registry_short:
                p = maya_obj_registry_short[self.parent_name]
                p.add(self)
            else:
                raise Exception("Parent for object not yet defined")

        maya_obj_registry[self.get_full_name()] = self
        maya_obj_registry_short[self.name] = self
39
40

    def _parse_first_maya_line(l):
41
        """ Extracts name and parent from maya file """
42
43
44
        name = parent = None
        vals = l.split()
        for i in range(len(vals)):
45
46
47
48
            if vals[i] == "-n": 
                name = MayaObj._strip_maya_string( vals[i+1] )
            if vals[i] == "-p":
                parent = MayaObj._strip_maya_string( vals[i+1] )
49
        return name,parent
50
51
52
53
54
55
56
57
58
59
60
61
62
63

    def _parse_maya_lines(self, maya_lines):
        """ Sets position and orientation from attributes in vHelix maya files """
        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]])
                elif vals[1] in ('".r"','".rsrr"'):
                    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 )
                    self.orientation = Rz.dot(Ry.dot(Rx))
cmaffeo2's avatar
cmaffeo2 committed
64
65
66
67
68
                elif vals[1] == '".lb"':
                    try:
                        self.sequence = MayaObj.seq_mapping[vals[2][0]]
                    except:
                        pass
69

70
    def _strip_maya_string(string):
71
        return re.match('"\|?([a-zA-Z_0-9|]+)";?$', string).group(1)
72

73
74
75
76
77
78
    def add(self, obj):
        self.children[obj.name] = obj
        obj.parent = self

    def get_position(self):
        if self.parent is not None:
79
            return self.parent.get_orientation().dot(self.position) + self.parent.get_position()
80
81
82
83
84
        else:
            return self.position

    def get_orientation(self):
        if self.parent is not None:
85
            return self.parent.get_orientation().dot(self.orientation)
86
87
88
89
90
91
92
93
        else:
            return self.orientation

    def get_full_name(self):
        ret = self.name
        if self.parent is not None:
            ret = '{}|{}'.format(self.parent.get_full_name(), ret)
        return ret
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

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

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

cmaffeo2's avatar
cmaffeo2 committed
132
class MayaBase(MayaObj):
133
    def __init__(self, maya_lines):
cmaffeo2's avatar
cmaffeo2 committed
134
        self.sequence = '?'
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        MayaObj.__init__(self, maya_lines)
        self._parse_maya_lines(maya_lines)        
        self.basepair = None
        self.end3 = None
        self.end5 = None

    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)
156

157
158
    def __repr__(self):
        return "Base {} {:.2f}".format(self.name, self.position[2])
159

cmaffeo2's avatar
cmaffeo2 committed
160
def parse_maya_file(maya_file):
161
162
163
164
    """ Function to parse vHelix maya ascii file, extracting useful information about base """

    objects = []
    objects_by_full_name = dict()
cmaffeo2's avatar
cmaffeo2 committed
165
166
167

    helices = dict()
    bases = []
168
    aim_constraints = []
169
    connections = []
cmaffeo2's avatar
cmaffeo2 committed
170

171
    """ Parse .ma file """
cmaffeo2's avatar
cmaffeo2 committed
172
173
174
    with open(maya_file) as fh:
        linesBuffer = []
        for line in fh:
175
176
177
            words = line.split()
            if words[0] == "createNode":

cmaffeo2's avatar
cmaffeo2 committed
178
179
180
181
                ## Create a new object
                if len(linesBuffer) > 0:
                    objType = linesBuffer[0].split()[1]
                    if objType == "vHelix":
182
183
                        h = MayaObj( linesBuffer )
                        helices[h.get_full_name()] = h
cmaffeo2's avatar
cmaffeo2 committed
184
185
                    elif objType == "HelixBase":
                        bases.append( MayaBase( linesBuffer ) )
186
187
188
189
190
                    elif objType == "transform":
                        o = MayaObj( linesBuffer )
                        objects.append(o)
                    elif objType == "aimConstraint":
                        aim_constraints.append( MayaObj( 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)

199
    """ Parse connections """
200
201
202
203
204
205
206
    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
207

208
    """ Assign base connectivity """
209
210
211
212
213
214
215
216
    for c in connections:
        h1 = c.helix1
        bn1 = c.base1
        s1 = c.suff1
        h2 = c.helix2
        bn2 = c.base2
        s2 = c.suff2
        
217
218
        b1 = helices[h1].children[bn1]
        b2 = helices[h2].children[bn2]
219
220
221
222
        if s1 == "lb" and s2 == "lb":
            b1.add_basepair(b2)
        elif s1 == "fw" and s2 == "bw":
            b1.add_end5(b2)
223
224
        elif s1 == "bw" and s2 == "fw":
            b1.add_end3(b2)
225
        else:
226
            raise Exception("Unrecognized connection %s %s" % (s1,s2))
cmaffeo2's avatar
cmaffeo2 committed
227

228
229
    """ Find local basis of each base so that
    model_from_basepair_stack_3prime can determine stacks """
230

231
232
233
234
235
    ref_below_position = np.array((5.11, -0.4, -3.34))
    ref_stack_position = np.array((-1.5, -4.7, 3.34))
    ref_bp_position = np.array((0.042, -16.98, 0.0))
    ref_bp_below_position = np.array((4.2, -12.9, 3.34))
    ref_bp_stack_position = np.array((-4.42, -14.46, -3.34))
236

237
238
239
240
241
    ref_positions = np.array((ref_below_position,
                              ref_stack_position,
                              ref_bp_position,
                              ref_bp_below_position,
                              ref_bp_stack_position))
242

243
    shift = np.array((1.19, -3.52, 0.08))
244

245
246
    if DEBUG:
        write_pdb_psf(bases, 'before')
cmaffeo2's avatar
cmaffeo2 committed
247

248
    def update_geometry(b,beads):
cmaffeo2's avatar
cmaffeo2 committed
249
        
250
251
        """ Update the position and orientation of bead b using the 5
        neighboring bases in "beads" as a reference """
252
        
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        def _get_beads_index(i):
            j=0
            while j <= i:
                if beads[j] is None:
                    i+=1
                j+=1
            return j-1

        """ Get position and filter out beads == None """
        r0 = b.get_position()
        tmp_ref_positions = ref_positions[[b2 is not None for b2 in beads]]
        positions = np.array([b2.get_position()-r0 for b2 in beads if b2 is not None])
        dist = np.linalg.norm(positions,axis=-1)

        """ Discard neighboring bases that are very far from current base """
        if np.any(dist > 20) and len(positions) > 3:
            i = _get_beads_index(np.argmax(dist))
            beads[i] = None
            return update_geometry(b, beads)

        """ Find transform that fits neighborhood around base to a reference """
        R,comB,comA = minimizeRmsd(positions, tmp_ref_positions)

        dr = np.array([R.T.dot(p-comB)-p0+comA for p,p0 in zip(positions, tmp_ref_positions)])
        dr2 = (dr**2).sum(axis=-1)
        msd = dr2.sum()/len(positions)

        """ Discard neighboring bases that are very far from their expected position """
        if msd > 17 and len(positions) > 3:
            # print("WARNING:",b,msd)
            # print(dr2)
            i = _get_beads_index(np.argmax(dr2))
            beads[i] = None
            return update_geometry(b,beads)
287
        
288
289
290
291
292
293
294
295
        """ Assign position and orientation to base """
        if b.parent is None:
            o_parent = np.eye(3)
        else:
            o_parent = b.parent.get_orientation()

        b.orientation = o_parent.T.dot( R )
        b.position = b.position + b.orientation.dot( shift )
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
    """ Loop over bases updating their positions and orientations """
    for b in bases:
        if b.basepair is not None:
            beads = [b.end5, b.end3, b.basepair, b.basepair.end5, b.basepair.end3]
            update_geometry(b,beads)

    if DEBUG:
        write_pdb_psf(bases, 'after')

    return bases


def write_pdb_psf(bases, prefix):

    """ Function for debugging .ma parser """

    from ..model.arbdmodel import ArbdModel,ParticleType, PointParticle, Group
    from ..model.nonbonded import HarmonicBond

    types = {c:ParticleType(c.upper(), mass=1,radius=1) for c in "bfuxyzs"}
    bond=HarmonicBond(k=0,r0=1)

    base_to_particle = dict()
    seg = []
    count=0
    for b in bases:
        r = b.get_position()
        o = b.get_orientation()

        type_character = 'f' if re.match('.*forw.*',b.name) else 'b' if re.match('.*bac.*',b.name) else 'u'
        name = type_character + b.name.split("_")[-1]
        p = PointParticle( name=name, type_=types[type_character], position = r )
        particles = [p]
        base_to_particle[b] = p
        for axis,c in zip(np.eye(3),"xyz"):
            particles.append( PointParticle( name=c, type_=types[c],
                                        position = r + o.T.dot(axis) ))

        ref_stack_position = np.array((-2.41851735, -0.259761333, 3.39999978))
        particles.append( PointParticle( name=c, type_=types[c],
                                         position = r + o.dot(ref_stack_position) ))

        res = Group(name='r'+str(count), children=particles, resid=count)
        count+=1
        for p2 in particles[1:]:
            res.add_bond(p,p2, bond=bond)
        seg.append(res)

    model = ArbdModel(seg)
    for b in bases:
        if b.end3 is not None:
            model.add_bond( base_to_particle[b], base_to_particle[b.end3], bond=bond )
        if b.basepair is not None:
            model.add_bond( base_to_particle[b], base_to_particle[b.basepair], bond=bond )
351

352
353
354
355
    model._countParticleTypes()
    model._updateParticleOrder()
    model.writePsf(prefix+'.psf')
    model.writePdb(prefix+'.pdb')
356
357


358
def convert_maya_bases_to_segment_model(maya_bases, **model_parameters):
359

360
    """ Converts bases parse from .ma file into mrdna an model """
361

362
    from .segmentmodel_from_lists import model_from_basepair_stack_3prime
363

364
365
366
367
    base_to_index = {b:i for i,b in zip(range(len(maya_bases)),maya_bases)}
    coord = np.zeros((len(maya_bases),3))
    orientation = np.zeros((len(maya_bases),3,3))
    bp,end3 = [-np.ones((len(maya_bases)),dtype=np.int) for i in range(2)]
cmaffeo2's avatar
cmaffeo2 committed
368
369
370
371
    sequence = []

    seq_map = {'A':'T', 'C':'G', 'T':'A', 'G':'C'}

372

373
374
375
376
377
378
379
    for i,b in zip(range(len(maya_bases)),maya_bases):
        coord[i] = b.get_position()
        orientation[i] = b.get_orientation().T
        if b.basepair is not None:
            bp[i] = base_to_index[ b.basepair ]
        if b.end3 is not None:
            end3[i] = base_to_index[ b.end3 ]
cmaffeo2's avatar
cmaffeo2 committed
380

cmaffeo2's avatar
cmaffeo2 committed
381
382
383
384
385
386
387
388
389
        try:
            s = b.sequence
            if s == '?':
                s = seq_map[s.basepair.sequence]
        except:
            s = 'T'
        sequence.append( s )
    if 'sequence' not in model_parameters:
        model_parameters['sequence'] = sequence
390
391
392
    return model_from_basepair_stack_3prime(coord, bp, None, end3,
                                            orientation=orientation,
                                            **model_parameters)