from pymol import cmd from pymol.cgo import * def vvsub( a, b ): return a[0]-b[0], a[1]-b[1], a[2]-b[2] def oprod( a, b ): return a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2], a[0]*b[1]-a[1]*b[0] def triangle( x, y, z, rgb = [ 1, 1, 1 ] ): n = oprod( vvsub( y, x ), vvsub( z, x ) ) return [ COLOR, rgb[0], rgb[1], rgb[2], NORMAL, n[0], n[1], n[2], VERTEX, x[0], x[1], x[2], VERTEX, y[0], y[1], y[2], VERTEX, z[0], z[1], z[2] ] def square( k, l, m, n, rgb = [ 1, 1, 1 ] ): return triangle( k, l, m, rgb ) + triangle( m, n, k, rgb ) def cube( rgbt = [ 1, 1, 1, 0 ], name="cube" ): a = [ -1, -1, -1 ] b = [ -1, -1, 1 ] c = [ -1, 1, -1 ] d = [ -1, 1, 1 ] e = [ 1, -1, -1 ] f = [ 1, -1, 1 ] g = [ 1, 1, -1 ] h = [ 1, 1, 1 ] obj = [ BEGIN, TRIANGLES, ALPHA, rgbt[3] ] obj.extend( square( a, b, d, c, rgbt ) ) obj.extend( square( a, e, f, b, rgbt ) ) obj.extend( square( a, c, g, e, rgbt ) ) obj.extend( square( c, d, h, g, rgbt ) ) obj.extend( square( b, f, h, d, rgbt ) ) obj.extend( square( e, g, h, f, rgbt ) ) obj.extend( [ END ] ) cmd.load_cgo( obj, name ) def rect( p=[0,0,0], q=[1,1,1], rgbt = [ 1, 1, 1, 0 ], name="rect" ): u = [ min(p[0],q[0]), min(p[1],q[1]), min(p[2],q[2]) ] v = [ max(p[0],q[0]), max(p[1],q[1]), max(p[2],q[2]) ] a = u b = [ u[0], u[1], v[2] ] c = [ u[0], v[1], u[2] ] d = [ u[0], v[1], v[2] ] e = [ v[0], u[1], u[2] ] f = [ v[0], u[1], v[2] ] g = [ v[0], v[1], u[2] ] h = v obj = [ BEGIN, TRIANGLES, ALPHA, rgbt[3], ] obj.extend( square( a, b, d, c, rgbt ) ) obj.extend( square( a, e, f, b, rgbt ) ) obj.extend( square( a, c, g, e, rgbt ) ) obj.extend( square( c, d, h, g, rgbt ) ) obj.extend( square( b, f, h, d, rgbt ) ) obj.extend( square( e, g, h, f, rgbt ) ) obj.extend( [ END ] ) cmd.load_cgo( obj, name )