py2Dmol 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (9) hide show
  1. py2Dmol/__init__.py +1 -0
  2. py2Dmol/resources/__init__.py +0 -0
  3. py2Dmol/resources/pseudo_3D_viewer.html +1002 -0
  4. py2Dmol/viewer.py +198 -0
  5. py2dmol-1.0.0.dist-info/METADATA +29 -0
  6. py2dmol-1.0.0.dist-info/RECORD +9 -0
  7. py2dmol-1.0.0.dist-info/WHEEL +5 -0
  8. py2dmol-1.0.0.dist-info/licenses/LICENSE +8 -0
  9. py2dmol-1.0.0.dist-info/top_level.txt +1 -0
py2Dmol/__init__.py ADDED
@@ -0,0 +1 @@
1
+ from .viewer import view
py2Dmol/resources/__init__.py ADDED
File without changes
py2Dmol/resources/pseudo_3D_viewer.html ADDED
@@ -0,0 +1,1002 @@
1
+ <!DOCTYPE html>
2
+ <html lang="en">
3
+ <head>
4
+ <meta charset="UTF-8">
5
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
6
+ <title>Protein Pseudo-3D Viewer</title>
7
+ <style>
8
+ * { box-sizing: border-box; }
9
+ /* Remove all margins for a clean embed */
10
+ body {
11
+ margin: 0;
12
+ padding: 0;
13
+ font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
14
+ background: #fff;
15
+ text-align: left; /* Aligns the container to the left */
16
+ }
17
+ #mainContainer {
18
+ display: inline-block; /* Makes the container wrap content */
19
+ text-align: left;
20
+ }
21
+ #canvasContainer {
22
+ /* This is the stylized "box" */
23
+ display: inline-block;
24
+ position: relative;
25
+ border: 1px solid #ddd;
26
+ border-radius: 12px;
27
+ overflow: hidden;
28
+ box-shadow: 0 4px 12px rgba(0, 0, 0, 0.08);
29
+ }
30
+ #canvas {
31
+ background: #ffffff;
32
+ cursor: grab;
33
+ display: block;
34
+ }
35
+ #canvas:active {
36
+ cursor: grabbing;
37
+ }
38
+
39
+ /* Style for the floating dropdown */
40
+ #colorSelect {
41
+ position: absolute;
42
+ top: 10px;
43
+ right: 10px;
44
+ z-index: 10;
45
+ font-size: 12px;
46
+ padding: 4px 8px;
47
+ border: 1px solid #ccc;
48
+ border-radius: 4px;
49
+ background-color: rgba(255, 255, 255, 0.8);
50
+ -webkit-appearance: none;
51
+ -moz-appearance: none;
52
+ appearance: none;
53
+ cursor: pointer;
54
+ background-image: url('data:image/svg+xml;charset=US-ASCII,%3Csvg%2Mxmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20width%3D%22292.4%22%20height%3D%22292.4%22%3E%3Cpath%20fill%3D%22%23444444%22%20d%3D%22M287%2069.4a17.6%2017.6%200%200%200-13-5.4H18.4c-5%200-9.3%201.8-12.9%205.4A17.6%2017.6%200%200%200%200%2082.2c0%205%201.8%209.3%205.4%2012.9l128%20127.9c3.6%203.6%207.8%205.4%2012.8%205.4s9.2-1.8%2012.8-5.4L287%2095c3.6-3.6%205.4-7.8%205.4-12.8%200-5-1.8-9.2-5.4-12.8z%22%2F%3E%3C%2Fsvg%3E');
55
+ background-repeat: no-repeat;
56
+ background-position: right 8px top 50%;
57
+ background-size: 8px auto;
58
+ padding-right: 28px;
59
+ }
60
+ #colorSelect:focus {
61
+ outline: none;
62
+ border-color: #007bff;
63
+ }
64
+
65
+ /* NEW: Style for options container */
66
+ #optionsContainer {
67
+ position: absolute;
68
+ top: 10px;
69
+ left: 10px;
70
+ z-index: 10;
71
+ font-size: 12px;
72
+ color: #333;
73
+ background-color: rgba(255, 255, 255, 0.8);
74
+ padding: 4px 8px;
75
+ border-radius: 4px;
76
+ display: flex;
77
+ flex-direction: column; /* Stack controls vertically */
78
+ align-items: flex-start; /* Align to the left */
79
+ gap: 4px;
80
+ }
81
+ .toggle-item { /* NEW: Wrapper for label/input pairs */
82
+ display: flex;
83
+ align-items: center;
84
+ gap: 4px;
85
+ }
86
+ .toggle-item label { cursor: pointer; white-space: nowrap; }
87
+ .toggle-item input[type="checkbox"] { cursor: pointer; margin: 0; }
88
+ .toggle-item input[type="range"] {
89
+ cursor: pointer;
90
+ width: 80px; /* Short slider for linewidth */
91
+ }
92
+
93
+
94
+ /* Animation Controls Container */
95
+ #controlsContainer {
96
+ /* MODIFIED: Use flex for better alignment */
97
+ display: flex;
98
+ flex-wrap: nowrap;
99
+ align-items: center;
100
+ padding: 10px 10px 0 10px;
101
+ text-align: left;
102
+ width: 100%;
103
+ font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
104
+ font-size: 12px;
105
+ }
106
+ /* NEW: Trajectory Controls Container */
107
+ #trajectoryContainer {
108
+ padding: 5px 10px 10px 10px; /* Top padding 5px */
109
+ width: 100%;
110
+ font-size: 12px;
111
+ }
112
+ .controlButton { /* NEW: Shared button style */
113
+ padding: 5px 10px;
114
+ border-radius: 4px;
115
+ border: 1px solid #ccc;
116
+ background: #f0f0f0;
117
+ cursor: pointer;
118
+ min-width: 60px;
119
+ font-size: 12px;
120
+ vertical-align: middle;
121
+ margin-right: 5px;
122
+ flex-shrink: 0; /* Don't let buttons shrink */
123
+ }
124
+ .controlButton:disabled {
125
+ cursor: not-allowed;
126
+ background: #eee;
127
+ color: #999;
128
+ }
129
+ #frameSlider {
130
+ /* MODIFIED: Let it flex */
131
+ flex-grow: 1;
132
+ width: auto; /* Remove fixed/js width */
133
+ margin: 0 10px;
134
+ vertical-align: middle;
135
+ }
136
+ #frameCounter {
137
+ color: #333;
138
+ vertical-align: middle;
139
+ min-width: 80px;
140
+ display: inline-block;
141
+ flex-shrink: 0; /* Don't let it shrink */
142
+ }
143
+ #speedSelect { /* NEW */
144
+ font-size: 12px;
145
+ padding: 4px 8px;
146
+ border: 1px solid #ccc;
147
+ border-radius: 4px;
148
+ vertical-align: middle;
149
+ margin-left: 10px;
150
+ flex-shrink: 0; /* Don't let it shrink */
151
+ }
152
+ /* NEW: Trajectory Select */
153
+ #trajectorySelect {
154
+ font-size: 12px;
155
+ padding: 4px 8px;
156
+ border: 1px solid #ccc;
157
+ border-radius: 4px;
158
+ vertical-align: middle;
159
+ margin-left: 5px;
160
+ }
161
+ #trajectorySelect:disabled {
162
+ cursor: not-allowed;
163
+ background: #eee;
164
+ }
165
+ #trajectoryLabel {
166
+ vertical-align: middle;
167
+ }
168
+
169
+ </style>
170
+ </head>
171
+ <body>
172
+ <!-- Main container to hold viewer and controls -->
173
+ <div id="mainContainer">
174
+ <!-- The canvas and dropdown are now siblings in the container -->
175
+ <div id="canvasContainer">
176
+ <canvas id="canvas"></canvas>
177
+ <!-- MODIFIED: Options container -->
178
+ <div id="optionsContainer">
179
+ <div class="toggle-item">
180
+ <input type="checkbox" id="shadowCheckbox" checked>
181
+ <label for="shadowCheckbox">Shadow</label>
182
+ </div>
183
+ <!-- HTML TYPO FIX: class.toggle-item -> class="toggle-item" -->
184
+ <div class="toggle-item">
185
+ <input type="checkbox" id="rotationCheckbox">
186
+ <label for="rotationCheckbox">Rotate</label>
187
+ </div>
188
+ <div class="toggle-item">
189
+ <label for="lineWidthSlider">Line:</label>
190
+ <input type="range" id="lineWidthSlider" min="1" max="5" value="3" step="0.5">
191
+ </div>
192
+ </div>
193
+ <!-- Floating Dropdown Menu -->
194
+ <select id="colorSelect">
195
+ <option value="plddt">pLDDT</option>
196
+ <option value="rainbow">Rainbow</option>
197
+ <option value="chain">Chain</option>
198
+ </select>
199
+ </div>
200
+
201
+ <!-- Animation Controls -->
202
+ <div id="controlsContainer">
203
+ <button id="playButton" class="controlButton">Play</button>
204
+ <input type="range" id="frameSlider" min="0" max="0" value="0">
205
+ <span id="frameCounter">Frame: 0 / 0</span>
206
+ <!-- NEW: Speed Select -->
207
+ <select id="speedSelect">
208
+ <option value="100">1x</option>
209
+ <option value="50">2x</option>
210
+ <option value="25">4x</option>
211
+ </select>
212
+ </div>
213
+ <!-- NEW Trajectory Controls -->
214
+ <div id="trajectoryContainer">
215
+ <span id="trajectoryLabel">Trajectory:</span>
216
+ <select id="trajectorySelect">
217
+ <option value="default">0</option>
218
+ </select>
219
+ <!-- REMOVED Save Video Button -->
220
+ </div>
221
+ </div>
222
+
223
+ <!--
224
+ This single point will be replaced by the Python script
225
+ with viewer config AND initial data.
226
+ -->
227
+ <!-- DATA_INJECTION_POINT -->
228
+
229
+ <script>
230
+ // ============================================================================
231
+ // VECTOR MATH (Unchanged)
232
+ // ============================================================================
233
+ class Vec3 {
234
+ constructor(x, y, z) { this.x = x; this.y = y; this.z = z; }
235
+ add(v) { return new Vec3(this.x + v.x, this.y + v.y, this.z + v.z); }
236
+ sub(v) { return new Vec3(this.x - v.x, this.y - v.y, this.z - v.z); }
237
+ mul(s) { return new Vec3(this.x * s, this.y * s, this.z * s); }
238
+ dot(v) { return this.x * v.x + this.y * v.y + this.z * v.z; }
239
+ length() { return Math.sqrt(this.dot(this)); }
240
+ distanceTo(v) { return this.sub(v).length(); }
241
+ }
242
+ function rotationMatrixX(angle) { const c = Math.cos(angle), s = Math.sin(angle); return [[1,0,0], [0,c,-s], [0,s,c]]; }
243
+ function rotationMatrixY(angle) { const c = Math.cos(angle), s = Math.sin(angle); return [[c,0,s], [0,1,0], [-s,0,c]]; }
244
+ function multiplyMatrices(a, b) { const r = [[0,0,0],[0,0,0],[0,0,0]]; for (let i = 0; i < 3; i++) for (let j = 0; j < 3; j++) for (let k = 0; k < 3; k++) r[i][j] += a[i][k] * b[k][j]; return r; }
245
+ function applyMatrix(m, v) { return new Vec3(m[0][0]*v.x + m[0][1]*v.y + m[0][2]*v.z, m[1][0]*v.x + m[1][1]*v.y + m[1][2]*v.z, m[2][0]*v.x + m[2][1]*v.y + m[2][2]*v.z); }
246
+ function sigmoid(x) { return 1 / (1 + Math.exp(-x)); }
247
+
248
+ // ============================================================================
249
+ // COLOR UTILITIES (Unchanged)
250
+ // ============================================================================
251
+ const pymolColors = ["#33ff33","#00ffff","#ff33cc","#ffff00","#ff9999","#e5e5e5","#7f7fff","#ff7f00","#7fff7f","#199999","#ff007f","#ffdd5e","#8c3f99","#b2b2b2","#007fff","#c4b200","#8cb266","#00bfbf","#b27f7f","#fcd1a5","#ff7f7f","#ffbfdd","#7fffff","#ffff7f","#00ff7f","#337fcc","#d8337f","#bfff3f","#ff7fff","#d8d8ff","#3fffbf","#b78c4c","#339933","#66b2b2","#ba8c84","#84bf00","#b24c66","#7f7f7f","#3f3fa5","#a5512b"];
252
+ function hexToRgb(hex) { if (!hex || typeof hex !== 'string') { return {r: 128, g: 128, b: 128}; } const r = parseInt(hex.slice(1,3), 16); const g = parseInt(hex.slice(3,5), 16); const b = parseInt(hex.slice(5,7), 16); return {r, g, b}; }
253
+ function hsvToRgb(h, s, v) { const c = v * s; const x = c * (1 - Math.abs((h / 60) % 2 - 1)); const m = v - c; let r, g, b; if (h < 60) { r = c; g = x; b = 0; } else if (h < 120) { r = x; g = c; b = 0; } else if (h < 180) { r = 0; g = c; b = x; } else if (h < 240) { r = 0; g = x; b = c; } else if (h < 300) { r = x; g = 0; b = c; } else { r = c; g = 0; b = x; } return { r: Math.round((r + m) * 255), g: Math.round((g + m) * 255), b: Math.round((b + m) * 255) }; }
254
+ function getRainbowColor(value, min, max) { if (max - min < 1e-6) return hsvToRgb(0, 1.0, 1.0); let normalized = (value - min) / (max - min); normalized = Math.max(0, Math.min(1, normalized)); const hue = 240 * normalized; return hsvToRgb(hue, 1.0, 1.0); }
255
+ function getPlddtColor(plddt) { return getRainbowColor(plddt, 50, 90); }
256
+ function getChainColor(chainIndex) { if (chainIndex < 0) chainIndex = 0; return hexToRgb(pymolColors[chainIndex % pymolColors.length]); }
257
+
258
+ // ============================================================================
259
+ // DEMO STRUCTURE GENERATOR (Unchanged)
260
+ // ============================================================================
261
+ function generateProteinCurve(n) { const coords = []; const plddts = []; const chains = []; const atomTypes = []; let angle = 0, z = 0; const numProtein = n - 10; for (let i = 0; i < n; i++) { const t = i / n; chains.push(i < n / 2 ? 'A' : 'B'); if (i >= numProtein) { atomTypes.push('L'); plddts.push(75.0); const protEnd = coords[numProtein - 1] || new Vec3(0,0,0); const lt = (i - numProtein) / 10.0; const lx = protEnd.x + 5 + lt * 5 * Math.cos(lt * 20); const ly = protEnd.y + 5 + lt * 5 * Math.sin(lt * 20); const lz = protEnd.z - 3 + lt * 3; coords.push(new Vec3(lx, ly, lz)); } else { atomTypes.push('P'); plddts.push((Math.sin(t * Math.PI * 6) * 0.5 + 0.5) * 50 + 40); const freq = 2 + Math.sin(t * Math.PI * 2) * 1.5; angle += 0.3 * freq; const radius = 15 + 10 * Math.sin(t * Math.PI * 4); const x = radius * Math.cos(angle); const y = radius * Math.sin(angle); z += 0.5 + 0.3 * Math.sin(t * Math.PI * 8); coords.push(new Vec3(x, y, z)); } } return {coords, plddts, chains, atomTypes}; }
262
+
263
+ // ============================================================================
264
+ // PSEUDO-3D RENDERER (Refactored)
265
+ // ============================================================================
266
+ class Pseudo3DRenderer {
267
+ constructor(canvas) {
268
+ this.canvas = canvas;
269
+ this.ctx = canvas.getContext('2d');
270
+
271
+ // Current render state
272
+ this.coords = [];
273
+ this.plddts = [];
274
+ this.chains = [];
275
+ this.atomTypes = [];
276
+
277
+ // Viewer state
278
+ this.colorMode = 'plddt';
279
+ this.rotationMatrix = [[1,0,0],[0,1,0],[0,0,1]];
280
+ this.zoom = 1.0;
281
+ this.lineWidth = 3.0;
282
+ this.shadowIntensity = 0.95;
283
+ this.shadowEnabled = true;
284
+
285
+ // Performance
286
+ this.chainRainbowScales = {};
287
+
288
+ // --- REFACTORED: Animation & State ---
289
+ this.trajectoriesData = { "default": { maxExtent: 0, frames: [], globalCenterSum: new Vec3(0,0,0), totalAtoms: 0 } };
290
+ this.currentTrajectoryName = "default";
291
+ this.currentFrame = -1;
292
+
293
+ // Playback
294
+ this.isPlaying = false;
295
+ this.animationSpeed = 100; // ms per frame
296
+ this.lastFrameAdvanceTime = 0;
297
+
298
+ // Interaction
299
+ this.isDragging = false;
300
+ this.autoRotate = false;
301
+
302
+ // Inertia
303
+ this.spinVelocityX = 0;
304
+ this.spinVelocityY = 0;
305
+ this.lastDragTime = 0;
306
+ this.lastDragX = 0;
307
+ this.lastDragY = 0;
308
+
309
+ // Track slider interaction
310
+ this.isSliderDragging = false;
311
+
312
+ // UI elements
313
+ this.playButton = null;
314
+ // REMOVED: this.saveVideoButton = null;
315
+ this.frameSlider = null;
316
+ this.frameCounter = null;
317
+ this.trajectorySelect = null;
318
+ this.controlsContainer = null;
319
+ this.speedSelect = null;
320
+ this.rotationCheckbox = null;
321
+ this.lineWidthSlider = null;
322
+
323
+ // REMOVED: Video recording properties
324
+ // this.mediaRecorder = null;
325
+ // this.recordedChunks = [];
326
+ // this.isRecording = false;
327
+
328
+ this.setupInteraction();
329
+ }
330
+
331
+ setupInteraction() {
332
+ // --- REFACTORED: Added inertia logic ---
333
+ this.canvas.addEventListener('mousedown', (e) => {
334
+ // Only start dragging if we clicked directly on the canvas
335
+ if (e.target !== this.canvas) return;
336
+
337
+ this.isDragging = true;
338
+ this.spinVelocityX = 0;
339
+ this.spinVelocityY = 0;
340
+ this.lastDragX = e.clientX;
341
+ this.lastDragY = e.clientY;
342
+ this.lastDragTime = performance.now();
343
+ if (this.autoRotate) {
344
+ this.autoRotate = false;
345
+ if (this.rotationCheckbox) this.rotationCheckbox.checked = false;
346
+ }
347
+ });
348
+
349
+ window.addEventListener('mousemove', (e) => {
350
+ if (!this.isDragging) return; // REMOVED: isRecording check
351
+
352
+ // SLIDER FIX: Clear isDragging and stop if over a control element
353
+ if (e.target.tagName === 'INPUT' || e.target.tagName === 'SELECT' || e.target.tagName === 'BUTTON') {
354
+ this.isDragging = false;
355
+ return;
356
+ }
357
+
358
+ const now = performance.now();
359
+ const timeDelta = now - this.lastDragTime;
360
+
361
+ const dx = e.clientX - this.lastDragX;
362
+ const dy = e.clientY - this.lastDragY;
363
+
364
+ if (dy !== 0) { const rot = rotationMatrixX(dy * 0.01); this.rotationMatrix = multiplyMatrices(rot, this.rotationMatrix); }
365
+ if (dx !== 0) { const rot = rotationMatrixY(dx * 0.01); this.rotationMatrix = multiplyMatrices(rot, this.rotationMatrix); }
366
+
367
+ // Store velocity for inertia
368
+ if (timeDelta > 0) {
369
+ // Weighted average to smooth out jerky movements
370
+ const smoothing = 0.5;
371
+ this.spinVelocityX = (this.spinVelocityX * (1-smoothing)) + ((dx / timeDelta * 20) * smoothing);
372
+ this.spinVelocityY = (this.spinVelocityY * (1-smoothing)) + ((dy / timeDelta * 20) * smoothing);
373
+ }
374
+
375
+ this.lastDragX = e.clientX;
376
+ this.lastDragY = e.clientY;
377
+ this.lastDragTime = now;
378
+
379
+ this.render(); // Render immediately while dragging
380
+ });
381
+
382
+ window.addEventListener('mouseup', () => {
383
+ this.isDragging = false;
384
+ const now = performance.now();
385
+ const timeDelta = now - this.lastDragTime;
386
+
387
+ if (timeDelta > 100) { // If drag was too slow, or just a click
388
+ this.spinVelocityX = 0;
389
+ this.spinVelocityY = 0;
390
+ }
391
+ // Else, the velocity from the last mousemove is used by the animate loop
392
+ });
393
+
394
+ this.canvas.addEventListener('wheel', (e) => {
395
+ // REMOVED: isRecording check
396
+ e.preventDefault();
397
+ this.zoom *= (1 - e.deltaY * 0.001);
398
+ this.zoom = Math.max(0.1, Math.min(5, this.zoom));
399
+ this.render();
400
+ }, { passive: false });
401
+ }
402
+
403
+ // Set UI controls from main script
404
+ setUIControls(controlsContainer, playButton, frameSlider, frameCounter, trajectorySelect, speedSelect, rotationCheckbox, lineWidthSlider) { // MODIFIED: removed saveVideoButton
405
+ this.controlsContainer = controlsContainer;
406
+ this.playButton = playButton;
407
+ // REMOVED: this.saveVideoButton = saveVideoButton;
408
+ this.frameSlider = frameSlider;
409
+ this.frameCounter = frameCounter;
410
+ this.trajectorySelect = trajectorySelect;
411
+ this.speedSelect = speedSelect;
412
+ this.rotationCheckbox = rotationCheckbox;
413
+ this.lineWidthSlider = lineWidthSlider;
414
+
415
+ this.lineWidth = parseFloat(this.lineWidthSlider.value);
416
+
417
+ // --- BIND ALL EVENT LISTENERS HERE ---
418
+ this.playButton.addEventListener('click', () => {
419
+ this.togglePlay();
420
+ });
421
+
422
+ this.trajectorySelect.addEventListener('change', () => {
423
+ this.stopAnimation();
424
+ this.currentTrajectoryName = this.trajectorySelect.value;
425
+ this.setFrame(0);
426
+ });
427
+
428
+ this.speedSelect.addEventListener('change', (e) => {
429
+ this.animationSpeed = parseInt(e.target.value);
430
+ });
431
+
432
+ this.rotationCheckbox.addEventListener('change', (e) => {
433
+ this.autoRotate = e.target.checked;
434
+ // Stop inertia if user clicks auto-rotate
435
+ this.spinVelocityX = 0;
436
+ this.spinVelocityY = 0;
437
+ });
438
+
439
+ this.lineWidthSlider.addEventListener('input', (e) => {
440
+ this.lineWidth = parseFloat(e.target.value);
441
+ if (!this.isPlaying) { // REMOVED: isRecording check
442
+ this.render();
443
+ }
444
+ });
445
+
446
+ // --- SLIDER FIX: Prevent canvas drag from interfering with slider ---
447
+ const handleSliderChange = (e) => {
448
+ this.stopAnimation();
449
+ this.setFrame(parseInt(e.target.value));
450
+ };
451
+
452
+ // Track when user is interacting with slider
453
+ this.frameSlider.addEventListener('mousedown', (e) => {
454
+ this.isDragging = false;
455
+ this.isSliderDragging = true;
456
+ e.stopPropagation();
457
+ });
458
+
459
+ this.frameSlider.addEventListener('mouseup', (e) => {
460
+ this.isSliderDragging = false;
461
+ });
462
+
463
+ // Also clear on window mouseup in case user releases outside slider
464
+ window.addEventListener('mouseup', () => {
465
+ this.isSliderDragging = false;
466
+ });
467
+
468
+ this.frameSlider.addEventListener('input', handleSliderChange);
469
+ this.frameSlider.addEventListener('change', handleSliderChange);
470
+
471
+ // Also prevent canvas drag when interacting with other controls
472
+ const allControls = [this.playButton, this.trajectorySelect, this.speedSelect,
473
+ this.rotationCheckbox, this.lineWidthSlider];
474
+ allControls.forEach(control => {
475
+ if (control) {
476
+ control.addEventListener('mousedown', (e) => {
477
+ this.isDragging = false;
478
+ e.stopPropagation();
479
+ });
480
+ }
481
+ });
482
+ }
483
+
484
+ // Add a new trajectory
485
+ addTrajectory(name) {
486
+ this.stopAnimation();
487
+ this.trajectoriesData[name] = { maxExtent: 0, frames: [], globalCenterSum: new Vec3(0,0,0), totalAtoms: 0 };
488
+ this.currentTrajectoryName = name;
489
+ this.currentFrame = -1;
490
+
491
+ const option = document.createElement('option');
492
+ option.value = name;
493
+ option.textContent = name;
494
+ this.trajectorySelect.appendChild(option);
495
+ this.trajectorySelect.value = name;
496
+
497
+ this.updateUIControls();
498
+ }
499
+
500
+ // Add a frame (data is raw parsed JSON)
501
+ addFrame(data) {
502
+ if (!this.currentTrajectoryName) {
503
+ this.currentTrajectoryName = "default";
504
+ }
505
+ if (!this.trajectoriesData[this.currentTrajectoryName]) {
506
+ this.trajectoriesData[this.currentTrajectoryName] = { maxExtent: 0, frames: [], globalCenterSum: new Vec3(0,0,0), totalAtoms: 0 };
507
+ }
508
+
509
+ const trajectory = this.trajectoriesData[this.currentTrajectoryName];
510
+ trajectory.frames.push(data);
511
+
512
+ // --- Update global center sum and count ---
513
+ let frameSum = new Vec3(0,0,0);
514
+ let frameAtoms = 0;
515
+ if (data && data.coords) {
516
+ frameAtoms = data.coords.length;
517
+ for (const c of data.coords) {
518
+ frameSum = frameSum.add(new Vec3(c[0], c[1], c[2]));
519
+ }
520
+ trajectory.globalCenterSum = trajectory.globalCenterSum.add(frameSum);
521
+ trajectory.totalAtoms += frameAtoms;
522
+ }
523
+
524
+ const globalCenter = (trajectory.totalAtoms > 0) ? trajectory.globalCenterSum.mul(1 / trajectory.totalAtoms) : new Vec3(0,0,0);
525
+
526
+ // --- Recalculate maxExtent for *all* frames using the *new* global center ---
527
+ let maxDistSq = 0;
528
+ for (const frame of trajectory.frames) {
529
+ if (frame && frame.coords) {
530
+ for (const c of frame.coords) {
531
+ const coordVec = new Vec3(c[0], c[1], c[2]);
532
+ const centeredCoord = coordVec.sub(globalCenter);
533
+ const distSq = centeredCoord.dot(centeredCoord);
534
+ if (distSq > maxDistSq) maxDistSq = distSq;
535
+ }
536
+ }
537
+ }
538
+ trajectory.maxExtent = Math.sqrt(maxDistSq);
539
+
540
+ if (!this.isPlaying) { // REMOVED: isRecording check
541
+ this.setFrame(trajectory.frames.length - 1);
542
+ }
543
+ this.updateUIControls();
544
+ }
545
+
546
+ // Set the current frame and render it
547
+ setFrame(frameIndex) {
548
+ frameIndex = parseInt(frameIndex);
549
+ if (!this.currentTrajectoryName) return;
550
+
551
+ const trajectory = this.trajectoriesData[this.currentTrajectoryName];
552
+ if (!trajectory || frameIndex < 0 || frameIndex >= trajectory.frames.length) {
553
+ this.currentFrame = -1;
554
+ this.ctx.fillStyle = '#ffffff';
555
+ this.ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
556
+ this.updateUIControls();
557
+ return;
558
+ }
559
+
560
+ this.currentFrame = frameIndex;
561
+ const data = trajectory.frames[frameIndex];
562
+ this._loadDataIntoRenderer(data); // This calls render()
563
+ this.updateUIControls(); // Update slider value
564
+ }
565
+
566
+ // Update UI element states (e.g., disabled)
567
+ setUIEnabled(enabled) {
568
+ this.playButton.disabled = !enabled;
569
+ // REMOVED: this.saveVideoButton.disabled = !enabled;
570
+ this.frameSlider.disabled = !enabled;
571
+ this.trajectorySelect.disabled = !enabled;
572
+ this.speedSelect.disabled = !enabled;
573
+ this.rotationCheckbox.disabled = !enabled;
574
+ this.lineWidthSlider.disabled = !enabled;
575
+ this.canvas.style.cursor = enabled ? 'grab' : 'wait';
576
+ }
577
+
578
+ // Update the text/slider values
579
+ updateUIControls() {
580
+ if (!this.playButton) return;
581
+
582
+ const trajectory = this.trajectoriesData[this.currentTrajectoryName];
583
+ const total = trajectory ? trajectory.frames.length : 0;
584
+ const current = Math.max(0, this.currentFrame) + 1;
585
+
586
+ if (total <= 1) { // MODIFIED: Simplified show/hide
587
+ this.controlsContainer.style.display = 'none';
588
+ } else {
589
+ this.controlsContainer.style.display = 'flex';
590
+ }
591
+
592
+ this.frameSlider.max = Math.max(0, total - 1);
593
+
594
+ // CRITICAL FIX: Don't update slider value while user is dragging it!
595
+ if (!this.isSliderDragging) {
596
+ this.frameSlider.value = this.currentFrame;
597
+ }
598
+
599
+ this.frameCounter.textContent = `Frame: ${total > 0 ? current : 0} / ${total}`;
600
+
601
+ // MODIFIED: Simplified text logic
602
+ this.playButton.textContent = this.isPlaying ? 'Pause' : 'Play';
603
+ }
604
+
605
+ // Toggle play/pause
606
+ togglePlay() {
607
+ // REMOVED: isRecording check
608
+ if (this.isPlaying) {
609
+ this.stopAnimation();
610
+ } else {
611
+ this.startAnimation();
612
+ }
613
+ }
614
+
615
+ // --- REFACTORED: Start playback ---
616
+ startAnimation() {
617
+ // REMOVED: isRecording check
618
+ const trajectory = this.trajectoriesData[this.currentTrajectoryName];
619
+ if (!trajectory || trajectory.frames.length < 2) return;
620
+
621
+ this.isPlaying = true;
622
+ this.lastFrameAdvanceTime = performance.now(); // Set start time
623
+ this.updateUIControls();
624
+ }
625
+
626
+ // --- REFACTORED: Stop playback ---
627
+ stopAnimation() {
628
+ this.isPlaying = false;
629
+ this.updateUIControls();
630
+ }
631
+
632
+ // REMOVED: saveAnimationAsVideo method
633
+
634
+ // Load data into renderer
635
+ _loadDataIntoRenderer(data) {
636
+ try {
637
+ if (data && data.coords && data.coords.length > 0) {
638
+ const coords = data.coords.map(c => new Vec3(c[0], c[1], c[2]));
639
+ const plddts = data.plddts || [];
640
+ const chains = data.chains || [];
641
+ const atomTypes = data.atom_types || [];
642
+ this.setCoords(coords, plddts, chains, atomTypes);
643
+ }
644
+ } catch (e) {
645
+ console.error("Failed to load data into renderer:", e);
646
+ }
647
+ this.render();
648
+ }
649
+
650
+ // Set current coordinates
651
+ setCoords(coords, plddts = [], chains = [], atomTypes = []) {
652
+ this.coords = coords;
653
+ this.plddts = plddts;
654
+ this.chains = chains;
655
+ this.atomTypes = atomTypes;
656
+
657
+ if (this.plddts.length !== this.coords.length) { this.plddts = Array(this.coords.length).fill(50.0); }
658
+ if (this.chains.length !== this.coords.length) { this.chains = Array(this.coords.length).fill('A'); }
659
+ if (this.atomTypes.length !== this.coords.length) { this.atomTypes = Array(this.coords.length).fill('P'); }
660
+
661
+ this.chainRainbowScales = {};
662
+ for (let i = 0; i < this.atomTypes.length; i++) {
663
+ if (this.atomTypes[i] === 'P') {
664
+ const chainId = this.chains[i] || 'A';
665
+ if (!this.chainRainbowScales[chainId]) { this.chainRainbowScales[chainId] = { min: Infinity, max: -Infinity }; }
666
+ const colorIndex = this.coords.length - 1 - i;
667
+ const scale = this.chainRainbowScales[chainId];
668
+ scale.min = Math.min(scale.min, colorIndex);
669
+ scale.max = Math.max(scale.max, colorIndex);
670
+ }
671
+ }
672
+ }
673
+
674
+ // --- RENDER (Core drawing logic) ---
675
+ render() {
676
+ const startTime = performance.now();
677
+ this.ctx.fillStyle = '#ffffff';
678
+ this.ctx.fillRect(0, 0, this.canvas.width, this.canvas.height);
679
+ if (this.coords.length === 0) return;
680
+
681
+ const trajectory = this.trajectoriesData[this.currentTrajectoryName];
682
+ const globalCenter = (trajectory && trajectory.totalAtoms > 0) ? trajectory.globalCenterSum.mul(1 / trajectory.totalAtoms) : new Vec3(0,0,0);
683
+
684
+ const rotated = this.coords.map(v => applyMatrix(this.rotationMatrix, v.sub(globalCenter)));
685
+ const segments = [];
686
+ const proteinChainbreak = 5.0;
687
+ const ligandBondCutoff = 2.0;
688
+ let firstProteinIndex = -1;
689
+ let lastProteinIndex = -1;
690
+ const ligandIndices = [];
691
+
692
+ for (let i = 0; i < rotated.length; i++) {
693
+ if (this.atomTypes[i] === 'L') {
694
+ ligandIndices.push(i);
695
+ continue;
696
+ }
697
+ if (firstProteinIndex === -1) { firstProteinIndex = i; }
698
+ lastProteinIndex = i;
699
+ if (i < rotated.length - 1) {
700
+ const type1 = this.atomTypes[i];
701
+ const type2 = this.atomTypes[i+1];
702
+ if (type1 === 'P' && type2 === 'P') {
703
+ const start = rotated[i];
704
+ const end = rotated[i+1];
705
+ const dist = start.distanceTo(end);
706
+ if (dist < proteinChainbreak) {
707
+ segments.push({ start, end, mid: start.add(end).mul(0.5), length: dist, colorIndex: rotated.length - 1 - i, origIndex: i, chainId: this.chains[i] || 'A' });
708
+ }
709
+ }
710
+ }
711
+ }
712
+
713
+ if (firstProteinIndex !== -1 && lastProteinIndex !== -1 && firstProteinIndex !== lastProteinIndex) {
714
+ const firstProteinChainId = this.chains[firstProteinIndex] || 'A';
715
+ const lastProteinChainId = this.chains[lastProteinIndex] || 'A';
716
+ if (firstProteinChainId === lastProteinChainId) {
717
+ const start = rotated[firstProteinIndex];
718
+ const end = rotated[lastProteinIndex];
719
+ const dist = start.distanceTo(end);
720
+ if (dist < proteinChainbreak) {
721
+ segments.push({ start, end, mid: start.add(end).mul(0.5), length: dist, colorIndex: this.chainRainbowScales[firstProteinChainId]?.min || 0, origIndex: firstProteinIndex, chainId: firstProteinChainId });
722
+ }
723
+ }
724
+ }
725
+
726
+ for (let i = 0; i < ligandIndices.length; i++) {
727
+ for (let j = i + 1; j < ligandIndices.length; j++) {
728
+ const idx1 = ligandIndices[i];
729
+ const idx2 = ligandIndices[j];
730
+ const start = rotated[idx1];
731
+ const end = rotated[idx2];
732
+ const dist = start.distanceTo(end);
733
+ if (dist < ligandBondCutoff) {
734
+ segments.push({ start, end, mid: start.add(end).mul(0.5), length: dist, colorIndex: 0, origIndex: idx1, chainId: this.chains[idx1] || 'A' });
735
+ }
736
+ }
737
+ }
738
+
739
+ if (segments.length === 0) return;
740
+
741
+ const grey = {r: 128, g: 128, b: 128};
742
+ const colors = segments.map(seg => {
743
+ const i = seg.origIndex;
744
+ const type = this.atomTypes[i];
745
+ if (type === 'L') {
746
+ if (this.colorMode === 'plddt') {
747
+ const plddt1 = (this.plddts[i] !== null && this.plddts[i] !== undefined) ? this.plddts[i] : 50;
748
+ return getPlddtColor(plddt1);
749
+ }
750
+ return grey;
751
+ }
752
+ if (this.colorMode === 'plddt') {
753
+ const plddt1 = (this.plddts[i] !== null && this.plddts[i] !== undefined) ? this.plddts[i] : 50;
754
+ const plddt2_idx = (seg.origIndex + 1 < this.coords.length) ? seg.origIndex + 1 : seg.origIndex;
755
+ const plddt2 = (this.plddts[plddt2_idx] !== null && this.plddts[plddt2_idx] !== undefined) ? this.plddts[plddt2_idx] : 50;
756
+ return getPlddtColor((plddt1 + plddt2) / 2);
757
+ }
758
+ else if (this.colorMode === 'chain') {
759
+ if (this.chains.length === 0) return getChainColor(0);
760
+ const chainId = this.chains[i] || 'A';
761
+ const uniqueChains = [...new Set(this.chains)];
762
+ const chainIndex = uniqueChains.indexOf(chainId);
763
+ return getChainColor(chainIndex >= 0 ? chainIndex : 0);
764
+ }
765
+ else {
766
+ const scale = this.chainRainbowScales[seg.chainId];
767
+ if (scale) { return getRainbowColor(seg.colorIndex, scale.min, scale.max); }
768
+ else { return grey; }
769
+ }
770
+ });
771
+
772
+ const zValues = segments.map(s => (s.start.z + s.end.z) / 2);
773
+ const zMin = Math.min(...zValues);
774
+ const zMax = Math.max(...zValues);
775
+ const zNorm = zValues.map(z => zMax - zMin > 1e-6 ? (z - zMin) / (zMax - zMin) : 0);
776
+ const n = segments.length;
777
+ const shadows = new Float32Array(n);
778
+ const tints = new Float32Array(n);
779
+
780
+ if (this.shadowEnabled) {
781
+ for (let i = 0; i < n; i++) {
782
+ let shadowSum = 0; let maxTint = 0;
783
+ const seg1 = segments[i];
784
+ for (let j = 0; j < n; j++) {
785
+ if (i === j) continue;
786
+ const seg2 = segments[j];
787
+ if (zValues[i] >= zValues[j]) continue;
788
+ const avgLen = (seg1.length + seg2.length) / 2;
789
+ const shadow_cutoff = avgLen * 2.0;
790
+ const tint_cutoff = avgLen / 2.0;
791
+ const max_cutoff = shadow_cutoff + 10.0;
792
+ const dx = seg1.mid.x - seg2.mid.x; if (Math.abs(dx) > max_cutoff) continue;
793
+ const dy = seg1.mid.y - seg2.mid.y; if (Math.abs(dy) > max_cutoff) continue;
794
+ const dist2D = Math.sqrt(dx*dx + dy*dy); if (dist2D > max_cutoff) continue;
795
+ const dist3D = seg1.mid.distanceTo(seg2.mid);
796
+ if (dist3D < max_cutoff) { shadowSum += sigmoid(shadow_cutoff - dist3D); }
797
+ if (dist2D < tint_cutoff + 10.0) { maxTint = Math.max(maxTint, sigmoid(tint_cutoff - dist2D)); }
798
+ }
799
+ shadows[i] = Math.pow(this.shadowIntensity, shadowSum);
800
+ tints[i] = 1 - maxTint;
801
+ }
802
+ } else {
803
+ shadows.fill(1.0);
804
+ tints.fill(1.0);
805
+ }
806
+
807
+ const order = Array.from({length: n}, (_, i) => i).sort((a, b) => zValues[a] - zValues[b]);
808
+
809
+ const maxExtent = (trajectory && trajectory.maxExtent > 0) ? trajectory.maxExtent : 30.0;
810
+ const dataRange = (maxExtent * 2) + this.lineWidth * 2;
811
+
812
+ const canvasSize = Math.min(this.canvas.width, this.canvas.height);
813
+ const scale = (canvasSize / dataRange) * this.zoom;
814
+ const pyFigWidthPixels = 480.0;
815
+ const pyPixelsPerData = pyFigWidthPixels / dataRange;
816
+ const baseLineWidthPixels = (this.lineWidth * pyPixelsPerData) * this.zoom;
817
+ const centerX = this.canvas.width / 2;
818
+ const centerY = this.canvas.height / 2;
819
+
820
+ for (const idx of order) {
821
+ const seg = segments[idx];
822
+ let {r, g, b} = colors[idx];
823
+ r /= 255; g /= 255; b /= 255;
824
+ const tintFactor = (0.50 * zNorm[idx] + 0.50 * tints[idx]) / 3;
825
+ r = r + (1 - r) * tintFactor;
826
+ g = g + (1 - g) * tintFactor;
827
+ b = b + (1 - b) * tintFactor;
828
+ const shadowFactor = 0.20 + 0.25 * zNorm[idx] + 0.55 * shadows[idx];
829
+ r *= shadowFactor; g *= shadowFactor; b *= shadowFactor;
830
+ const color = `rgb(${r*255|0},${g*255|0},${b*255|0})`;
831
+ const x1 = centerX + seg.start.x * scale; const y1 = centerY - seg.start.y * scale;
832
+ const x2 = centerX + seg.end.x * scale; const y2 = centerY - seg.end.y * scale;
833
+ this.ctx.beginPath();
834
+ this.ctx.moveTo(x1, y1);
835
+ this.ctx.lineTo(x2, y2);
836
+ this.ctx.strokeStyle = color;
837
+ const type = this.atomTypes[seg.origIndex];
838
+ this.ctx.lineWidth = (type === 'L') ? (baseLineWidthPixels / 2) : baseLineWidthPixels;
839
+ this.ctx.lineCap = 'round';
840
+ this.ctx.stroke();
841
+ }
842
+ }
843
+
844
+ // --- REFACTORED: Main animation loop ---
845
+ animate() {
846
+ const now = performance.now();
847
+ let needsRender = false;
848
+
849
+ // 1. Handle inertia/spin
850
+ if (!this.isDragging) { // REMOVED: isRecording check
851
+ if (Math.abs(this.spinVelocityX) > 0.0001) {
852
+ const rot = rotationMatrixY(this.spinVelocityX * 0.005);
853
+ this.rotationMatrix = multiplyMatrices(rot, this.rotationMatrix);
854
+ this.spinVelocityX *= 0.95; // Damping
855
+ needsRender = true;
856
+ } else {
857
+ this.spinVelocityX = 0;
858
+ }
859
+
860
+ if (Math.abs(this.spinVelocityY) > 0.0001) {
861
+ const rot = rotationMatrixX(this.spinVelocityY * 0.005);
862
+ this.rotationMatrix = multiplyMatrices(rot, this.rotationMatrix);
863
+ this.spinVelocityY *= 0.95; // Damping
864
+ needsRender = true;
865
+ } else {
866
+ this.spinVelocityY = 0;
867
+ }
868
+ }
869
+
870
+ // 2. Handle auto-rotate
871
+ if (this.autoRotate && !this.isDragging && this.spinVelocityX === 0 && this.spinVelocityY === 0) { // REMOVED: isRecording check
872
+ const rot = rotationMatrixY(0.005); // Constant rotation speed
873
+ this.rotationMatrix = multiplyMatrices(rot, this.rotationMatrix);
874
+ needsRender = true;
875
+ }
876
+
877
+ // 3. Handle frame playback
878
+ if (this.isPlaying) { // REMOVED: isRecording check
879
+ if (now - this.lastFrameAdvanceTime > this.animationSpeed) {
880
+ const trajectory = this.trajectoriesData[this.currentTrajectoryName];
881
+ if (trajectory && trajectory.frames.length > 0) {
882
+ let nextFrame = this.currentFrame + 1;
883
+ if (nextFrame >= trajectory.frames.length) {
884
+ nextFrame = 0;
885
+ }
886
+ this.setFrame(nextFrame); // This calls render()
887
+ this.lastFrameAdvanceTime = now;
888
+ needsRender = false; // setFrame() already called render()
889
+ } else {
890
+ this.stopAnimation();
891
+ }
892
+ }
893
+ }
894
+
895
+ // 4. Final render if needed
896
+ if (needsRender) {
897
+ this.render();
898
+ }
899
+
900
+ // 5. Loop
901
+ requestAnimationFrame(() => this.animate());
902
+ }
903
+ }
904
+
905
+ // ============================================================================
906
+ // MAIN APP & COLAB COMMUNICATION
907
+ // ============================================================================
908
+
909
+ // 1. Get config from Python
910
+ const config = window.viewerConfig || { size: [800, 600], color: "plddt" };
911
+
912
+ // 2. Setup Canvas
913
+ const canvas = document.getElementById('canvas');
914
+ canvas.width = config.size[0];
915
+ canvas.height = config.size[1];
916
+ document.getElementById('mainContainer').style.width = `${config.size[0]}px`;
917
+
918
+ // 3. Create renderer
919
+ window.renderer = new Pseudo3DRenderer(canvas);
920
+ window.renderer.colorMode = config.color;
921
+
922
+ // 4. Setup general controls
923
+ const colorSelect = document.getElementById('colorSelect');
924
+ colorSelect.value = config.color;
925
+ colorSelect.addEventListener('change', (e) => {
926
+ window.renderer.colorMode = e.target.value;
927
+ window.renderer.render();
928
+ });
929
+ const shadowCheckbox = document.getElementById('shadowCheckbox');
930
+ shadowCheckbox.checked = window.renderer.shadowEnabled;
931
+ shadowCheckbox.addEventListener('change', (e) => {
932
+ window.renderer.shadowEnabled = e.target.checked;
933
+ window.renderer.render();
934
+ });
935
+
936
+ // 5. Setup animation and trajectory controls
937
+ const controlsContainer = document.getElementById('controlsContainer');
938
+ const playButton = document.getElementById('playButton');
939
+ // REMOVED: const saveVideoButton = document.getElementById('saveVideoButton');
940
+ const frameSlider = document.getElementById('frameSlider');
941
+ const frameCounter = document.getElementById('frameCounter');
942
+ const trajectorySelect = document.getElementById('trajectorySelect');
943
+ const speedSelect = document.getElementById('speedSelect');
944
+ const rotationCheckbox = document.getElementById('rotationCheckbox');
945
+ const lineWidthSlider = document.getElementById('lineWidthSlider');
946
+
947
+
948
+ // Pass ALL controls to the renderer
949
+ window.renderer.setUIControls(
950
+ controlsContainer, playButton,
951
+ frameSlider, frameCounter, trajectorySelect,
952
+ speedSelect, rotationCheckbox, lineWidthSlider
953
+ ); // MODIFIED: removed saveVideoButton
954
+
955
+ // REFACTORED: All event listeners are now bound inside setUIControls
956
+ // playButton.addEventListener('click', ...);
957
+ // frameSlider.addEventListener('input', ...);
958
+ // frameSlider.addEventListener('change', ...);
959
+
960
+
961
+ // 6. Add function for Python to call (for new frames)
962
+ window.handlePythonUpdate = (jsonData) => {
963
+ try {
964
+ const data = JSON.parse(jsonData);
965
+ window.renderer.addFrame(data);
966
+ } catch (e) {
967
+ console.error("Failed to parse JSON from Python:", e);
968
+ }
969
+ };
970
+
971
+ // 7. Add function for Python to start a new trajectory
972
+ window.handlePythonNewTrajectory = (name) => {
973
+ window.renderer.addTrajectory(name);
974
+ };
975
+
976
+ // 8. Load initial data
977
+ try {
978
+ if (window.proteinData && window.proteinData.coords && window.proteinData.coords.length > 0) {
979
+ window.renderer.addFrame(window.proteinData);
980
+ } else {
981
+ const {coords, plddts, chains, atomTypes} = generateProteinCurve(100);
982
+ const demoData = {
983
+ coords: coords.map(c => [c.x, c.y, c.z]),
984
+ plddts: plddts,
985
+ chains: chains,
986
+ atom_types: atomTypes
987
+ };
988
+ window.renderer.addFrame(demoData);
989
+ }
990
+ } catch (error) {
991
+ console.error("Error loading initial data:", error);
992
+ const {coords, plddts, chains, atomTypes} = generateProteinCurve(100);
993
+ const demoData = { coords: coords.map(c => [c.x, c.v, c.z]), plddts: plddts, chains: chains, atom_types: atomTypes };
994
+ window.renderer.addFrame(demoData);
995
+ }
996
+
997
+ // 9. Start the main animation loop
998
+ window.renderer.animate();
999
+
1000
+ </script>
1001
+ </body>
1002
+ </html>
py2Dmol/viewer.py ADDED
@@ -0,0 +1,198 @@
1
+ import json
2
+ import numpy as np
3
+ try:
4
+ from google.colab import output
5
+ IS_COLAB = True
6
+ except ImportError:
7
+ IS_COLAB = False
8
+ from IPython.display import display, HTML, Javascript
9
+ import importlib.resources
10
+ from . import resources as py2dmol_resources
11
+ import gemmi
12
+
13
+ def kabsch(a, b, return_v=False):
14
+ """Computes the optimal rotation matrix for aligning a to b."""
15
+ ab = a.swapaxes(-1, -2) @ b
16
+ u, s, vh = np.linalg.svd(ab, full_matrices=False)
17
+ flip = np.linalg.det(u @ vh) < 0
18
+ flip_b = flip[..., None]
19
+ u_last_col_flipped = np.where(flip_b, -u[..., -1], u[..., -1])
20
+ u[..., -1] = u_last_col_flipped
21
+ R = u @ vh
22
+ return u if return_v else R
23
+
24
+ def align_a_to_b(a, b):
25
+ """Aligns coordinate set 'a' to 'b' using Kabsch algorithm."""
26
+ a_mean = a.mean(-2, keepdims=True)
27
+ a_cent = a - a_mean
28
+ b_mean = b.mean(-2, keepdims=True)
29
+ b_cent = b - b_mean
30
+ R = kabsch(a_cent, b_cent)
31
+ a_aligned = (a_cent @ R) + b_mean
32
+ return a_aligned
33
+
34
+ # --- view Class ---
35
+
36
+ class view:
37
+ def __init__(self, size=(500,500), color="rainbow"):
38
+ self.size = size
39
+ self.color = color
40
+ self._initial_data_loaded = False
41
+ self._coords = None
42
+ self._plddts = None
43
+ self._chains = None
44
+ self._atom_types = None
45
+ self._trajectory_counter = 0 # NEW: Counter for trials
46
+
47
+ def _serialize_data(self):
48
+ """Serializes the current coordinate state to JSON."""
49
+ payload = {
50
+ "coords": self._coords.tolist(),
51
+ "plddts": self._plddts.tolist(),
52
+ "chains": list(self._chains),
53
+ "atom_types": list(self._atom_types)
54
+ }
55
+ return json.dumps(payload)
56
+
57
+ def _update(self, coords, plddts=None, chains=None, atom_types=None):
58
+ """Updates the internal state with new data, aligning coords."""
59
+ if self._coords is None:
60
+ self._coords = coords
61
+ else:
62
+ # Align new coords to old coords
63
+ # This prevents the structure from "jumping" if the center moves
64
+ self._coords = align_a_to_b(coords, self._coords)
65
+
66
+ # Set defaults if not provided
67
+ if self._plddts is None: self._plddts = np.full(self._coords.shape[0], 50.0) # Default to 50
68
+ if self._chains is None: self._chains = ["A"] * self._coords.shape[0]
69
+ if self._atom_types is None: self._atom_types = ["P"] * self._coords.shape[0]
70
+
71
+ # Update with new data if provided
72
+ if plddts is not None: self._plddts = plddts
73
+ if chains is not None: self._chains = chains
74
+ if atom_types is not None: self._atom_types = atom_types
75
+
76
+ # Ensure all arrays have the same length as coords
77
+ if len(self._plddts) != len(self._coords):
78
+ print(f"Warning: pLDDT length mismatch. Resetting to default.")
79
+ self._plddts = np.full(self._coords.shape[0], 50.0)
80
+ if len(self._chains) != len(self._coords):
81
+ print(f"Warning: Chains length mismatch. Resetting to default.")
82
+ self._chains = ["A"] * self._coords.shape[0]
83
+ if len(self._atom_types) != len(self._coords):
84
+ print(f"Warning: Atom types length mismatch. Resetting to default.")
85
+ self._atom_types = ["P"] * self._coords.shape[0]
86
+
87
+ def clear(self):
88
+ """MODIFIED: Clears the Python state and tells the JS viewer to start a new trajectory."""
89
+ # Generate a name for the new trajectory
90
+ trajectory_name = f"{self._trajectory_counter}"
91
+ self._trajectory_counter += 1
92
+
93
+ # Clear Python-side coordinates
94
+ self._coords = None
95
+ self._plddts = None
96
+ self._chains = None
97
+ self._atom_types = None
98
+
99
+ # NEW: Tell JS to create and switch to this new trajectory
100
+ if self._initial_data_loaded:
101
+ js_code = f"window.handlePythonNewTrajectory('{trajectory_name}');"
102
+ if IS_COLAB:
103
+ try:
104
+ output.eval_js(js_code, ignore_result=True)
105
+ except Exception as e:
106
+ pass
107
+ else:
108
+ display(Javascript(js_code))
109
+
110
+ def display(self, initial_coords, initial_plddts=None, initial_chains=None, initial_atom_types=None):
111
+ """Displays the viewer with initial data."""
112
+ self._update(initial_coords, initial_plddts, initial_chains, initial_atom_types)
113
+
114
+ try:
115
+ with importlib.resources.open_text(py2dmol_resources, 'pseudo_3D_viewer.html') as f:
116
+ html_template = f.read()
117
+ except FileNotFoundError:
118
+ print("Error: Could not find the HTML template file.")
119
+ return
120
+
121
+ viewer_config = {
122
+ "size": self.size,
123
+ "color": self.color
124
+ }
125
+ config_script = f"""
126
+ <script id="viewer-config">
127
+ window.viewerConfig = {json.dumps(viewer_config)};
128
+ </script>
129
+ """
130
+ # The data script now provides the *first frame* for the "Initial" trajectory
131
+ data_script = f"""
132
+ <script id="protein-data">
133
+ window.proteinData = {self._serialize_data()};
134
+ </script>
135
+ """
136
+ self._initial_data_loaded = True
137
+
138
+ injection_scripts = config_script + "\n" + data_script
139
+ final_html = html_template.replace("<!-- DATA_INJECTION_POINT -->", injection_scripts)
140
+ display(HTML(final_html))
141
+
142
+ def add(self, coords, plddts=None, chains=None, atom_types=None):
143
+ """Sends a new frame of data to the JavaScript viewer."""
144
+ if self._initial_data_loaded:
145
+ self._update(coords, plddts, chains, atom_types)
146
+ json_data = self._serialize_data()
147
+ js_code = f"window.handlePythonUpdate(`{json_data}`);"
148
+ if IS_COLAB:
149
+ output.eval_js(js_code, ignore_result=True)
150
+ else:
151
+ display(Javascript(js_code))
152
+ else:
153
+ # If display() was never called, call it now
154
+ self.display(coords, plddts, chains, atom_types)
155
+
156
+ update_data = add
157
+
158
+ def from_pdb(self, filepath, chains=None):
159
+ """Loads a structure from a PDB or CIF file and updates the viewer."""
160
+ structure = gemmi.read_structure(filepath)
161
+ self.clear()
162
+
163
+ for model in structure:
164
+ coords = []
165
+ plddts = []
166
+ atom_chains = []
167
+ atom_types = []
168
+
169
+ for chain in model:
170
+ if chains is None or chain.name in chains:
171
+ for residue in chain:
172
+ if residue.name == 'HOH':
173
+ continue
174
+
175
+ is_protein = gemmi.find_tabulated_residue(residue.name).is_amino_acid()
176
+
177
+ if is_protein:
178
+ if 'CA' in residue:
179
+ atom = residue['CA'][0]
180
+ coords.append(atom.pos.tolist())
181
+ plddts.append(atom.b_iso)
182
+ atom_chains.append(chain.name)
183
+ atom_types.append('P')
184
+ else: # Ligand
185
+ for atom in residue:
186
+ if atom.element.name != 'H':
187
+ coords.append(atom.pos.tolist())
188
+ plddts.append(atom.b_iso)
189
+ atom_chains.append(chain.name)
190
+ atom_types.append('L')
191
+
192
+ if coords:
193
+ coords = np.array(coords)
194
+ plddts = np.array(plddts)
195
+ if self._initial_data_loaded:
196
+ self.add(coords, plddts, atom_chains, atom_types)
197
+ else:
198
+ self.display(coords, plddts, atom_chains, atom_types)
py2dmol-1.0.0.dist-info/METADATA ADDED
@@ -0,0 +1,29 @@
1
+ Metadata-Version: 2.4
2
+ Name: py2Dmol
3
+ Version: 1.0.0
4
+ Summary: A Python library for visualizing protein structures in 2D.
5
+ Home-page: https://github.com/sokrypton/py2Dmol
6
+ Author: sokrypton
7
+ Author-email: so3@mit.edu
8
+ License: BEER-WARE
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.6
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE
14
+ Requires-Dist: numpy
15
+ Requires-Dist: ipython
16
+ Requires-Dist: gemmi
17
+ Dynamic: author
18
+ Dynamic: author-email
19
+ Dynamic: classifier
20
+ Dynamic: description
21
+ Dynamic: description-content-type
22
+ Dynamic: home-page
23
+ Dynamic: license
24
+ Dynamic: license-file
25
+ Dynamic: requires-dist
26
+ Dynamic: requires-python
27
+ Dynamic: summary
28
+
29
+ A Python library for visualizing protein structures in 2D.
py2dmol-1.0.0.dist-info/RECORD ADDED
@@ -0,0 +1,9 @@
1
+ py2Dmol/__init__.py,sha256=2WEbB5I-lcKuR1qMvJn4t1Z2iFHLdCnGRlRDRmLeJE8,25
2
+ py2Dmol/viewer.py,sha256=wv-57ejrQKYTn-ZVAGZPA6R9u3bJujmb-LYpScStJOQ,7632
3
+ py2Dmol/resources/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
+ py2Dmol/resources/pseudo_3D_viewer.html,sha256=1GFzNUGIxNkSQI0c_lCyXWNkyMuvRx-dcFsKztdV7Po,47913
5
+ py2dmol-1.0.0.dist-info/licenses/LICENSE,sha256=ysi-yHz7oTnMzDxfESZ40dxirc1eZeg2Wa6hFH81ZBQ,442
6
+ py2dmol-1.0.0.dist-info/METADATA,sha256=Vm5t0jgn-uQn-exoDJtjRBcDpzy7xx3C168NYX2zhls,771
7
+ py2dmol-1.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
8
+ py2dmol-1.0.0.dist-info/top_level.txt,sha256=uptjIQ7j7_exwdPSDocJkY7eNyY_YeA_Cw6KdwKrKi0,8
9
+ py2dmol-1.0.0.dist-info/RECORD,,
py2dmol-1.0.0.dist-info/WHEEL ADDED
@@ -0,0 +1,5 @@
1
+ Wheel-Version: 1.0
2
+ Generator: setuptools (80.9.0)
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
5
+
py2dmol-1.0.0.dist-info/licenses/LICENSE ADDED
@@ -0,0 +1,8 @@
1
+ /*
2
+ * ----------------------------------------------------------------------------
3
+ * "THE BEER-WARE LICENSE" (Revision 42):
4
+ * <so3@mit.edu> wrote this file. As long as you retain this notice you
5
+ * can do whatever you want with this stuff. If we meet some day, and you think
6
+ * this stuff is worth it, you can buy me a beer in return. Sergey Ovchinnikov
7
+ * ----------------------------------------------------------------------------
8
+ */
py2dmol-1.0.0.dist-info/top_level.txt ADDED
@@ -0,0 +1 @@
1
+ py2Dmol