vimol 0.1.0__tar.gz → 0.1.2__tar.gz

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 (50) hide show
  1. vimol-0.1.2/PKG-INFO +106 -0
  2. vimol-0.1.2/README.md +72 -0
  3. {vimol-0.1.0 → vimol-0.1.2}/pyproject.toml +8 -2
  4. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/__init__.py +1 -1
  5. vimol-0.1.2/src/vimol/_render_fast.py +418 -0
  6. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/app.py +58 -8
  7. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/editor.py +37 -28
  8. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/elements.py +23 -0
  9. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/gl_adapter.py +30 -5
  10. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/gl_render.py +109 -56
  11. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/input.py +22 -2
  12. vimol-0.1.2/src/vimol/kitty.py +654 -0
  13. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/parsers/xyz.py +1 -1
  14. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/render.py +206 -63
  15. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/scene.py +69 -13
  16. vimol-0.1.2/src/vimol/structures.py +398 -0
  17. vimol-0.1.2/src/vimol/theme.py +142 -0
  18. vimol-0.1.2/src/vimol/viewer.py +2477 -0
  19. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/widget.py +151 -27
  20. vimol-0.1.2/src/vimol.egg-info/PKG-INFO +106 -0
  21. {vimol-0.1.0 → vimol-0.1.2}/src/vimol.egg-info/SOURCES.txt +6 -0
  22. {vimol-0.1.0 → vimol-0.1.2}/src/vimol.egg-info/requires.txt +1 -2
  23. {vimol-0.1.0 → vimol-0.1.2}/tests/test_editing.py +705 -27
  24. {vimol-0.1.0 → vimol-0.1.2}/tests/test_gl_render.py +196 -0
  25. vimol-0.1.2/tests/test_render_fast.py +208 -0
  26. vimol-0.1.2/tests/test_scene_structures.py +357 -0
  27. vimol-0.1.2/tests/test_structures.py +421 -0
  28. vimol-0.1.2/tests/test_vimol.py +2632 -0
  29. vimol-0.1.0/PKG-INFO +0 -302
  30. vimol-0.1.0/README.md +0 -268
  31. vimol-0.1.0/src/vimol/kitty.py +0 -297
  32. vimol-0.1.0/src/vimol/viewer.py +0 -1159
  33. vimol-0.1.0/src/vimol.egg-info/PKG-INFO +0 -302
  34. vimol-0.1.0/tests/test_vimol.py +0 -552
  35. {vimol-0.1.0 → vimol-0.1.2}/LICENSE +0 -0
  36. {vimol-0.1.0 → vimol-0.1.2}/setup.cfg +0 -0
  37. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/__main__.py +0 -0
  38. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/arrows.py +0 -0
  39. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/bonds.py +0 -0
  40. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/camera.py +0 -0
  41. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/molecule.py +0 -0
  42. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/parsers/__init__.py +0 -0
  43. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/parsers/mol.py +0 -0
  44. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/parsers/pdb.py +0 -0
  45. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/periodic_table.py +0 -0
  46. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/py.typed +0 -0
  47. {vimol-0.1.0 → vimol-0.1.2}/src/vimol/templates.py +0 -0
  48. {vimol-0.1.0 → vimol-0.1.2}/src/vimol.egg-info/dependency_links.txt +0 -0
  49. {vimol-0.1.0 → vimol-0.1.2}/src/vimol.egg-info/entry_points.txt +0 -0
  50. {vimol-0.1.0 → vimol-0.1.2}/src/vimol.egg-info/top_level.txt +0 -0
vimol-0.1.2/PKG-INFO ADDED
@@ -0,0 +1,106 @@
1
+ Metadata-Version: 2.4
2
+ Name: vimol
3
+ Version: 0.1.2
4
+ Summary: A terminal molecular viewer and editor using the Kitty graphics protocol
5
+ Author: vimol contributors
6
+ License-Expression: MIT
7
+ Keywords: molecular,viewer,editor,terminal,kitty,chemistry,visualization,xyz,pdb,sdf,cli,tui
8
+ Classifier: Development Status :: 4 - Beta
9
+ Classifier: Environment :: Console
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Operating System :: OS Independent
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3 :: Only
15
+ Classifier: Programming Language :: Python :: 3.8
16
+ Classifier: Programming Language :: Python :: 3.9
17
+ Classifier: Programming Language :: Python :: 3.10
18
+ Classifier: Programming Language :: Python :: 3.11
19
+ Classifier: Programming Language :: Python :: 3.12
20
+ Classifier: Programming Language :: Python :: 3.13
21
+ Classifier: Topic :: Scientific/Engineering :: Chemistry
22
+ Classifier: Topic :: Scientific/Engineering :: Visualization
23
+ Classifier: Topic :: Multimedia :: Graphics :: Viewers
24
+ Classifier: Typing :: Typed
25
+ Requires-Python: >=3.8
26
+ Description-Content-Type: text/markdown
27
+ License-File: LICENSE
28
+ Requires-Dist: numpy>=1.20
29
+ Requires-Dist: moderngl>=5.6
30
+ Requires-Dist: numba>=0.59
31
+ Provides-Extra: test
32
+ Requires-Dist: pytest>=7; extra == "test"
33
+ Dynamic: license-file
34
+
35
+ # vimol
36
+
37
+ A molecular viewer and editor that runs **in your terminal**. Open an `.xyz`,
38
+ `.pdb`, `.mol`, or `.sdf` file and get a shaded, rotatable, editable 3D
39
+ structure right where you're already working — over SSH, in tmux, next to
40
+ your editor.
41
+
42
+ ```bash
43
+ pip install vimol
44
+ vimol molecule.xyz
45
+ ```
46
+
47
+ Requires Python ≥ 3.8 and a terminal that speaks the
48
+ [Kitty graphics protocol](https://sw.kovidgoyal.net/kitty/graphics-protocol/)
49
+ — **kitty, Ghostty, or WezTerm**. Rendering is pure software (numpy impostor
50
+ raycasting), with an OpenGL fast path used automatically where a GPU is
51
+ available. MIT licensed.
52
+
53
+ ## In the terminal
54
+
55
+ ![C60 spinning in the terminal, rendered by vimol](https://raw.githubusercontent.com/hanimustafa/vimol/main/docs/media/spin.gif)
56
+
57
+ `vimol` with no file opens this bundled C60 demo. Drag to rotate, scroll to
58
+ zoom, hover an atom to identify it, `m` to measure distances/angles/dihedrals
59
+ — with more than one structure loaded, a live column tracks it in a
60
+ comparison table next to the structure list, evaluated for every loaded
61
+ structure at once; move on to a different measurement (or a different frame)
62
+ and it locks in as its own column (click a column's `×` to remove it).
63
+ Editing is on by default: `a` to append (grow fragments, swap elements — the
64
+ status-bar pills pick the element and geometry), option-drag to draw bonds,
65
+ `x` to delete, `c` to relax clashes, `u` to undo, `s` to save. `?` lists
66
+ every binding.
67
+
68
+ ```bash
69
+ vimol traj.xyz --spin --style spacefill # spin a trajectory, space-filling
70
+ vimol protein.pdb --render out.png # batch still to PNG (works headless)
71
+ ```
72
+
73
+ ## Library usage
74
+
75
+ ```python
76
+ import vimol
77
+
78
+ mol = vimol.load("caffeine.sdf") # parse; bonds perceived if absent
79
+ vimol.view(mol, editable=True) # full-screen interactive viewer
80
+
81
+ scene = vimol.Scene(mol, 640, 480) # ...or drive the renderer yourself
82
+ scene.camera.orbit(30, -15) # rotate (degrees)
83
+ img = scene.render() # (H, W, 3) uint8 numpy array
84
+ os.write(1, scene.to_kitty()) # paint it into a Kitty terminal
85
+ scene.to_png("out.png") # or save a PNG (stdlib encoder, no Pillow)
86
+ ```
87
+
88
+ To embed the viewer in your own terminal app, keep the input loop and hand
89
+ only the events you want to the widget:
90
+
91
+ ```python
92
+ from vimol import MoleculeWidget, InputDecoder, MouseEvent
93
+
94
+ widget = MoleculeWidget(mol, width_px, height_px) # no terminal, no loop
95
+ decoder = InputDecoder(pixel=True)
96
+
97
+ for ev in decoder.feed(data): # bytes you read from the tty
98
+ if isinstance(ev, MouseEvent) and inside_my_region(ev):
99
+ widget.handle_mouse(ev, origin=(region_x_px, region_y_px))
100
+ else:
101
+ my_app_handles(ev)
102
+
103
+ os.write(1, widget.to_kitty(cols=region_cols, rows=region_rows))
104
+ ```
105
+
106
+ See `examples/embed_demo.py` for a complete host app.
vimol-0.1.2/README.md ADDED
@@ -0,0 +1,72 @@
1
+ # vimol
2
+
3
+ A molecular viewer and editor that runs **in your terminal**. Open an `.xyz`,
4
+ `.pdb`, `.mol`, or `.sdf` file and get a shaded, rotatable, editable 3D
5
+ structure right where you're already working — over SSH, in tmux, next to
6
+ your editor.
7
+
8
+ ```bash
9
+ pip install vimol
10
+ vimol molecule.xyz
11
+ ```
12
+
13
+ Requires Python ≥ 3.8 and a terminal that speaks the
14
+ [Kitty graphics protocol](https://sw.kovidgoyal.net/kitty/graphics-protocol/)
15
+ — **kitty, Ghostty, or WezTerm**. Rendering is pure software (numpy impostor
16
+ raycasting), with an OpenGL fast path used automatically where a GPU is
17
+ available. MIT licensed.
18
+
19
+ ## In the terminal
20
+
21
+ ![C60 spinning in the terminal, rendered by vimol](https://raw.githubusercontent.com/hanimustafa/vimol/main/docs/media/spin.gif)
22
+
23
+ `vimol` with no file opens this bundled C60 demo. Drag to rotate, scroll to
24
+ zoom, hover an atom to identify it, `m` to measure distances/angles/dihedrals
25
+ — with more than one structure loaded, a live column tracks it in a
26
+ comparison table next to the structure list, evaluated for every loaded
27
+ structure at once; move on to a different measurement (or a different frame)
28
+ and it locks in as its own column (click a column's `×` to remove it).
29
+ Editing is on by default: `a` to append (grow fragments, swap elements — the
30
+ status-bar pills pick the element and geometry), option-drag to draw bonds,
31
+ `x` to delete, `c` to relax clashes, `u` to undo, `s` to save. `?` lists
32
+ every binding.
33
+
34
+ ```bash
35
+ vimol traj.xyz --spin --style spacefill # spin a trajectory, space-filling
36
+ vimol protein.pdb --render out.png # batch still to PNG (works headless)
37
+ ```
38
+
39
+ ## Library usage
40
+
41
+ ```python
42
+ import vimol
43
+
44
+ mol = vimol.load("caffeine.sdf") # parse; bonds perceived if absent
45
+ vimol.view(mol, editable=True) # full-screen interactive viewer
46
+
47
+ scene = vimol.Scene(mol, 640, 480) # ...or drive the renderer yourself
48
+ scene.camera.orbit(30, -15) # rotate (degrees)
49
+ img = scene.render() # (H, W, 3) uint8 numpy array
50
+ os.write(1, scene.to_kitty()) # paint it into a Kitty terminal
51
+ scene.to_png("out.png") # or save a PNG (stdlib encoder, no Pillow)
52
+ ```
53
+
54
+ To embed the viewer in your own terminal app, keep the input loop and hand
55
+ only the events you want to the widget:
56
+
57
+ ```python
58
+ from vimol import MoleculeWidget, InputDecoder, MouseEvent
59
+
60
+ widget = MoleculeWidget(mol, width_px, height_px) # no terminal, no loop
61
+ decoder = InputDecoder(pixel=True)
62
+
63
+ for ev in decoder.feed(data): # bytes you read from the tty
64
+ if isinstance(ev, MouseEvent) and inside_my_region(ev):
65
+ widget.handle_mouse(ev, origin=(region_x_px, region_y_px))
66
+ else:
67
+ my_app_handles(ev)
68
+
69
+ os.write(1, widget.to_kitty(cols=region_cols, rows=region_rows))
70
+ ```
71
+
72
+ See `examples/embed_demo.py` for a complete host app.
@@ -34,10 +34,16 @@ classifiers = [
34
34
  "Topic :: Multimedia :: Graphics :: Viewers",
35
35
  "Typing :: Typed",
36
36
  ]
37
- dependencies = ["numpy>=1.20"]
37
+ # moderngl is a hard dependency (~120 KB of wheels): the GL backend is always
38
+ # available, and `backend="auto"` still falls back to the CPU raycaster at
39
+ # runtime wherever no OpenGL context can be created (headless boxes, clusters).
40
+ # numba compiles that CPU raycaster's inner loop (_render_fast) to machine
41
+ # code: ~10-25x faster CPU rendering, bit-identical output. render.py still
42
+ # tolerates numba being uninstalled/uncompilable (falls back to the pure-numpy
43
+ # path) since it's a best-effort JIT, not because it's meant to be omitted.
44
+ dependencies = ["numpy>=1.20", "moderngl>=5.6", "numba>=0.59"]
38
45
 
39
46
  [project.optional-dependencies]
40
- gl = ["moderngl>=5.6"]
41
47
  test = ["pytest>=7"]
42
48
 
43
49
  [project.scripts]
@@ -33,7 +33,7 @@ from . import elements
33
33
  from . import templates
34
34
  from . import editor
35
35
 
36
- __version__ = "0.1.0"
36
+ __version__ = "0.1.2"
37
37
 
38
38
  __all__ = [
39
39
  "Molecule", "VectorField", "Camera", "Renderer", "Style", "Scene", "MoleculeWidget",
@@ -0,0 +1,418 @@
1
+ """Numba-compiled inner loop for the CPU raycaster.
2
+
3
+ The numpy renderer (``render.py``) vectorizes per primitive over its
4
+ screen-space bounding box: every sphere/cylinder costs ~40 full-box array
5
+ passes, most of them over pixels that fail the z-test, plus thousands of
6
+ per-call overheads per frame. That is a structural ceiling -- measured
7
+ ~130 ms/frame for C60 at 2560x1440.
8
+
9
+ This module compiles the same impostor math to machine code (per-pixel loops
10
+ with early z-rejection *before* shading, the way PyMOL's C++ ray tracer and
11
+ every other compiled rasterizer do it), which measures 10-25x faster. numba
12
+ is a hard dependency (see pyproject.toml), but this module still degrades
13
+ gracefully if it's ever missing or fails to compile on some platform:
14
+ ``ready()`` stays False and ``render.py``'s numpy path runs unchanged.
15
+
16
+ BIT-PARITY CONTRACT: the kernels below mirror render.py's numpy operations
17
+ one-for-one -- same dtypes (float32 per-pixel math, float64 screen setup),
18
+ same operation order and association, same scalar casts (numpy converts a
19
+ python-float scalar to float32 when it meets a float32 array; typed numpy
20
+ float64 scalars promote instead -- both cases are reproduced explicitly).
21
+ ``tests/test_render_fast.py`` verifies the output is bit-identical to the
22
+ numpy path -- run it after touching either side of the contract.
23
+ """
24
+ from __future__ import annotations
25
+
26
+ import os
27
+ import threading
28
+
29
+ import numpy as np
30
+
31
+ # numba's disk cache (cache=True below) defaults to __pycache__ next to this
32
+ # source file -- churn a sync client (Dropbox) then uploads on every
33
+ # recompile. Redirect it to the user cache dir before numba is first
34
+ # imported (the env var is read at numba import; if some other package
35
+ # imported numba already, this is a harmless no-op and the default location
36
+ # applies). setdefault, so an explicit NUMBA_CACHE_DIR always wins.
37
+ os.environ.setdefault(
38
+ "NUMBA_CACHE_DIR",
39
+ os.path.join(os.path.expanduser("~"), ".cache", "vimol", "numba"))
40
+
41
+ try:
42
+ from numba import njit, prange
43
+ _HAVE_NUMBA = True
44
+ except Exception: # hard dependency, but degrade gracefully
45
+ _HAVE_NUMBA = False
46
+
47
+ def njit(*args, **kwargs): # stub so module import never fails
48
+ def wrap(fn):
49
+ return fn
50
+ return wrap if not (args and callable(args[0])) else args[0]
51
+
52
+ def prange(*args):
53
+ return range(*args)
54
+
55
+
56
+ # 0 = cold, 1 = compiling, 2 = ready, 3 = unavailable/failed
57
+ _state = 0
58
+ _lock = threading.Lock()
59
+ # Set once the one compile attempt (by whichever caller -- warm_async's
60
+ # thread or a direct warm_sync() caller -- got there first) finishes, success
61
+ # or failure. Callers that lost the race to start it wait on this instead of
62
+ # calling the compile routine themselves: numba's dispatcher is not safe to
63
+ # enter from two threads at once for the same not-yet-compiled function --
64
+ # observed reliably (SIGABRT/SIGSEGV) when a second, unsynchronized caller
65
+ # invoked the compile while warm_async's background thread was mid-compile.
66
+ _done = threading.Event()
67
+
68
+
69
+ def available() -> bool:
70
+ """numba importable on this machine (kernel may still be compiling)."""
71
+ return _HAVE_NUMBA
72
+
73
+
74
+ def ready() -> bool:
75
+ """Kernel compiled and safe to call (first call triggers compilation)."""
76
+ return _state == 2
77
+
78
+
79
+ def warm_async() -> None:
80
+ """Compile the kernel in a daemon thread so the first frames don't stall.
81
+
82
+ JIT compilation takes ~1-3 s once per process (cached on disk afterwards
83
+ where the package dir is writable). Frames rendered while this runs use
84
+ the numpy path; the renderer switches over the moment ready() flips.
85
+ """
86
+ if not _HAVE_NUMBA:
87
+ return
88
+ if not _claim():
89
+ return
90
+ threading.Thread(target=_run_once, daemon=True).start()
91
+
92
+
93
+ def warm_sync(timeout: float | None = None) -> bool:
94
+ """Force the kernel to be ready (or known unavailable) before returning.
95
+
96
+ If nobody has started the compile yet, this thread does it (blocking).
97
+ If warm_async's background thread already claimed it, this waits for
98
+ that thread to finish instead of compiling again -- see the note on
99
+ ``_done`` above for why a second concurrent attempt is unsafe. Returns
100
+ ``ready()``. Used by callers (tests, the benchmark scripts) that need a
101
+ definite answer and can't rely on warm_async's fire-and-forget timing.
102
+ """
103
+ if not _HAVE_NUMBA:
104
+ return False
105
+ if _claim():
106
+ _run_once()
107
+ else:
108
+ _done.wait(timeout)
109
+ return ready()
110
+
111
+
112
+ def _claim() -> bool:
113
+ """Atomically move state 0 -> 1. True for the one caller that must
114
+ actually run the compile; False for everyone else (already compiling,
115
+ ready, or failed)."""
116
+ global _state
117
+ with _lock:
118
+ if _state != 0:
119
+ return False
120
+ _state = 1
121
+ return True
122
+
123
+
124
+ def _run_once() -> None:
125
+ global _state
126
+ try:
127
+ # (8, 8, 4) rgba: (H, W, 3) and (H, W, 4) C-contiguous buffers share
128
+ # one numba type (uint8 3d C), so this single call compiles the only
129
+ # signature both opaque and transparent frames use.
130
+ color = np.zeros((8, 8, 4), np.uint8)
131
+ zbuf = np.full((8, 8), -np.inf, np.float32)
132
+ bands = np.array([[0, 8]], np.int64)
133
+ sx = np.array([4.0], np.float64)
134
+ sy = np.array([4.0], np.float64)
135
+ sz = np.array([0.0], np.float64)
136
+ radii = np.array([0.5], np.float64)
137
+ acol = np.array([[0.5, 0.5, 0.5]], np.float32)
138
+ aflat = np.zeros(1, np.uint8)
139
+ order = np.array([0], np.int64)
140
+ barr = np.zeros((0, 14), np.float64)
141
+ bcol = np.zeros((0, 6), np.float32)
142
+ bflat = np.zeros(0, np.uint8)
143
+ render_frame(color, zbuf, True, bands,
144
+ sx, sy, sz, radii, acol, aflat, order,
145
+ barr, bcol, bflat,
146
+ 10.0, 4.0, 4.0,
147
+ 0.3, 0.5, 0.7, -0.3, -0.5, 0.6, 0.4, 0.5, 0.7,
148
+ 0.28, 0.25, 0.55, 0.55, -1.0, 2.0, 5)
149
+ _state = 2
150
+ except Exception:
151
+ _state = 3
152
+ finally:
153
+ _done.set()
154
+
155
+
156
+ # ---------------------------------------------------------------------------
157
+ # compiled kernels
158
+ # ---------------------------------------------------------------------------
159
+ @njit(inline="always")
160
+ def _shade_pixel(color, y, x, alpha_on, depth, nx, ny, nz,
161
+ ar, ag, ab, flat,
162
+ l0, l1, l2, f0, f1, f2, hv0, hv1, hv2,
163
+ ambient, fill_w, spec_w, depth_cue, zmin, zspan, n_squares):
164
+ """Shade one winning fragment and write color/alpha/depth.
165
+
166
+ ``color`` is (H, W, 3) opaque or (H, W, 4) rgba; ``alpha_on`` means the
167
+ 4th channel exists and coverage is written there. Mirrors
168
+ render.shade_write's per-element op sequence (see module docstring).
169
+ nx/ny/nz/ar..ab/depth are float32; the light scalars are float64
170
+ python-side values that numpy would have cast to float32 at each use --
171
+ the casts here make that explicit.
172
+ """
173
+ if flat:
174
+ if depth_cue > 0.0:
175
+ fog = (depth - np.float32(zmin)) * np.float32(1.0 / zspan)
176
+ if fog < np.float32(0.0):
177
+ fog = np.float32(0.0)
178
+ if fog > np.float32(1.0):
179
+ fog = np.float32(1.0)
180
+ fog *= np.float32(depth_cue)
181
+ fog += np.float32(1.0 - depth_cue)
182
+ rr = ar * fog
183
+ rr *= np.float32(255.0)
184
+ gg = ag * fog
185
+ gg *= np.float32(255.0)
186
+ bb = ab * fog
187
+ bb *= np.float32(255.0)
188
+ else:
189
+ rr = min(ar, np.float32(1.0)) * np.float32(255.0)
190
+ gg = min(ag, np.float32(1.0)) * np.float32(255.0)
191
+ bb = min(ab, np.float32(1.0)) * np.float32(255.0)
192
+ else:
193
+ diff = nx * np.float32(l0) + ny * np.float32(l1) + nz * np.float32(l2)
194
+ if diff < np.float32(0.0):
195
+ diff = np.float32(0.0)
196
+ dfill = nx * np.float32(f0) + ny * np.float32(f1) + nz * np.float32(f2)
197
+ if dfill < np.float32(0.0):
198
+ dfill = np.float32(0.0)
199
+ dh = nx * np.float32(hv0) + ny * np.float32(hv1) + nz * np.float32(hv2)
200
+ if dh < np.float32(0.0):
201
+ dh = np.float32(0.0)
202
+ # _fast_pow for a power-of-two exponent: x*x then repeated squaring,
203
+ # identical rounding sequence to the numpy version.
204
+ spec = dh
205
+ for _ in range(n_squares):
206
+ spec *= spec
207
+ spec *= np.float32(spec_w)
208
+ diff += np.float32(ambient)
209
+ dfill *= np.float32(fill_w)
210
+ diff += dfill
211
+ if depth_cue > 0.0:
212
+ fog = (depth - np.float32(zmin)) * np.float32(1.0 / zspan)
213
+ if fog < np.float32(0.0):
214
+ fog = np.float32(0.0)
215
+ if fog > np.float32(1.0):
216
+ fog = np.float32(1.0)
217
+ fog *= np.float32(depth_cue)
218
+ fog += np.float32(1.0 - depth_cue)
219
+ diff *= fog
220
+ spec *= fog
221
+ rr = ar * diff
222
+ rr += spec
223
+ gg = ag * diff
224
+ gg += spec
225
+ bb = ab * diff
226
+ bb += spec
227
+ if rr > np.float32(1.0):
228
+ rr = np.float32(1.0)
229
+ if gg > np.float32(1.0):
230
+ gg = np.float32(1.0)
231
+ if bb > np.float32(1.0):
232
+ bb = np.float32(1.0)
233
+ rr *= np.float32(255.0)
234
+ gg *= np.float32(255.0)
235
+ bb *= np.float32(255.0)
236
+ # np.copyto(..., casting="unsafe") / setitem float->uint8 truncates.
237
+ color[y, x, 0] = np.uint8(rr)
238
+ color[y, x, 1] = np.uint8(gg)
239
+ color[y, x, 2] = np.uint8(bb)
240
+ if alpha_on:
241
+ color[y, x, 3] = 255
242
+ return
243
+
244
+
245
+ @njit(parallel=True, cache=True)
246
+ def _render_kernel(color, zbuf, alpha_on, bands,
247
+ sx, sy, sz, radii, acol, aflat, order,
248
+ barr, bcol, bflat,
249
+ zoom, ox_s, oy_s,
250
+ l0, l1, l2, f0, f1, f2, hv0, hv1, hv2,
251
+ ambient, fill_w, spec_w, depth_cue, zmin, zspan, n_squares):
252
+ """Raycast all bonds (list order) then all atoms (far-to-near order)
253
+ into their bands. Bands own disjoint rows, so prange is race-free; the
254
+ z-test makes per-pixel results order-independent up to the same ties the
255
+ numpy path resolves identically (same primitive order per band)."""
256
+ H, W = zbuf.shape
257
+ inv_zoom = 1.0 / zoom
258
+ iz32 = np.float32(1.0 / zoom)
259
+ niz32 = np.float32(-1.0 / zoom)
260
+ izz32 = np.float32(inv_zoom * inv_zoom)
261
+ n_bonds = barr.shape[0]
262
+
263
+ for band in prange(bands.shape[0]):
264
+ y_lo = bands[band, 0]
265
+ y_hi = bands[band, 1]
266
+
267
+ # -- bonds: capped-cylinder impostors (render._draw_cylinder_segment)
268
+ for k in range(n_bonds):
269
+ ax = barr[k, 0]
270
+ ay = barr[k, 1]
271
+ az = barr[k, 2]
272
+ ux = barr[k, 3]
273
+ uy = barr[k, 4]
274
+ uz = barr[k, 5]
275
+ L = barr[k, 6]
276
+ a2 = barr[k, 7]
277
+ radius = barr[k, 8]
278
+ sax = barr[k, 9]
279
+ say = barr[k, 10]
280
+ sbx = barr[k, 11]
281
+ sby = barr[k, 12]
282
+ srb = barr[k, 13]
283
+ if L < 1e-6 or radius <= 0.0:
284
+ continue
285
+ x0 = max(int(np.floor(min(sax, sbx) - srb)), 0)
286
+ x1 = min(int(np.ceil(max(sax, sbx) + srb)) + 1, W)
287
+ y0 = max(int(np.floor(min(say, sby) - srb)), y_lo)
288
+ y1 = min(int(np.ceil(max(say, sby) + srb)) + 1, y_hi)
289
+ if x0 >= x1 or y0 >= y1:
290
+ continue
291
+ ux32 = np.float32(ux)
292
+ uy32 = np.float32(uy)
293
+ uz32 = np.float32(uz)
294
+ b2c32 = np.float32(-2.0 * uz)
295
+ a2x4_32 = np.float32(a2 * 4.0)
296
+ half_a2_32 = np.float32(0.5 / a2) if abs(a2) >= 1e-9 else np.float32(0.0)
297
+ parallel = abs(a2) < 1e-9
298
+ inv_r32 = np.float32(1.0 / radius)
299
+ r2_32 = np.float32(radius * radius)
300
+ half_L32 = np.float32(0.5 * L)
301
+ L32 = np.float32(L)
302
+ az32 = np.float32(az)
303
+ ey_off32 = np.float32(oy_s / zoom - ay)
304
+ ox32 = np.float32(ox_s)
305
+ ax32 = np.float32(ax)
306
+ flat = bflat[k]
307
+ car, cag, cab = bcol[k, 0], bcol[k, 1], bcol[k, 2]
308
+ cbr, cbg, cbb = bcol[k, 3], bcol[k, 4], bcol[k, 5]
309
+ for y in range(y0, y1):
310
+ ey = np.float32(y) * niz32 + ey_off32
311
+ eyey = ey * ey
312
+ eyuy = ey * uy32
313
+ for x in range(x0, x1):
314
+ ex = (np.float32(x) - ox32) * iz32 - ax32
315
+ A0 = ex * ux32 + eyuy
316
+ c2 = ex * ex + eyey - A0 * A0 - r2_32
317
+ if parallel:
318
+ if c2 > np.float32(0.0):
319
+ continue
320
+ s = np.float32(0.0)
321
+ else:
322
+ b2 = A0 * b2c32
323
+ disc = b2 * b2 - a2x4_32 * c2
324
+ if disc < np.float32(0.0):
325
+ continue
326
+ sq = np.sqrt(disc)
327
+ s = (sq - b2) * half_a2_32
328
+ zview = s + az32
329
+ t = A0 + s * uz32
330
+ if t < np.float32(0.0) or t > L32:
331
+ continue
332
+ if zview <= zbuf[y, x]:
333
+ continue
334
+ nx = (ex - t * ux32) * inv_r32
335
+ ny = (ey - t * uy32) * inv_r32
336
+ nz = (s - t * uz32) * inv_r32
337
+ if t < half_L32:
338
+ ar, ag, ab = car, cag, cab
339
+ else:
340
+ ar, ag, ab = cbr, cbg, cbb
341
+ zbuf[y, x] = zview
342
+ _shade_pixel(color, y, x, alpha_on, zview, nx, ny, nz,
343
+ ar, ag, ab, flat,
344
+ l0, l1, l2, f0, f1, f2, hv0, hv1, hv2,
345
+ ambient, fill_w, spec_w, depth_cue, zmin, zspan,
346
+ n_squares)
347
+
348
+ # -- atoms: sphere impostors (render.draw_band's atom loop)
349
+ for oi in range(order.shape[0]):
350
+ a = order[oi]
351
+ r = radii[a]
352
+ if r <= 0.0:
353
+ continue
354
+ sr = r * zoom
355
+ if sr < 0.5:
356
+ sr = 0.5
357
+ cx = sx[a]
358
+ cy = sy[a]
359
+ cz = sz[a]
360
+ x0 = max(int(np.floor(cx - sr)), 0)
361
+ x1 = min(int(np.ceil(cx + sr)) + 1, W)
362
+ y0 = max(int(np.floor(cy - sr)), y_lo)
363
+ y1 = min(int(np.ceil(cy + sr)) + 1, y_hi)
364
+ if x0 >= x1 or y0 >= y1:
365
+ continue
366
+ cx32 = np.float32(cx)
367
+ cy32 = np.float32(cy)
368
+ cz32 = np.float32(cz)
369
+ sr2 = sr * sr # float64: d2 <= sr2 promotes
370
+ rr32 = np.float32(r * r)
371
+ inv_r = 1.0 / r
372
+ nxy32 = np.float32(inv_zoom * inv_r)
373
+ nz32 = np.float32(inv_r)
374
+ flat = aflat[a]
375
+ ar = acol[a, 0]
376
+ ag = acol[a, 1]
377
+ ab = acol[a, 2]
378
+ for y in range(y0, y1):
379
+ dyr = np.float32(y) - cy32
380
+ dyr2 = dyr * dyr
381
+ ny = dyr * np.float32(-inv_zoom * inv_r)
382
+ for x in range(x0, x1):
383
+ dxr = np.float32(x) - cx32
384
+ d2 = dxr * dxr + dyr2
385
+ # numpy compares float32 d2 against the float64 scalar
386
+ # sr*sr (typed numpy scalar promotes); mirror exactly.
387
+ if not (d2 <= sr2):
388
+ continue
389
+ h2 = rr32 - d2 * izz32
390
+ if h2 < np.float32(0.0):
391
+ h2 = np.float32(0.0)
392
+ hgt = np.sqrt(h2)
393
+ depth = hgt + cz32
394
+ if depth <= zbuf[y, x]:
395
+ continue
396
+ nx = dxr * nxy32
397
+ nzv = hgt * nz32
398
+ zbuf[y, x] = depth
399
+ _shade_pixel(color, y, x, alpha_on, depth,
400
+ nx, ny, nzv, ar, ag, ab, flat,
401
+ l0, l1, l2, f0, f1, f2, hv0, hv1, hv2,
402
+ ambient, fill_w, spec_w, depth_cue, zmin, zspan,
403
+ n_squares)
404
+
405
+
406
+ def render_frame(color, zbuf, alpha_on, bands,
407
+ sx, sy, sz, radii, acol, aflat, order,
408
+ barr, bcol, bflat,
409
+ zoom, ox_s, oy_s,
410
+ l0, l1, l2, f0, f1, f2, hv0, hv1, hv2,
411
+ ambient, fill_w, spec_w, depth_cue, zmin, zspan, n_squares):
412
+ """Thin wrapper so _run_once and render.py share one call signature."""
413
+ _render_kernel(color, zbuf, alpha_on, bands,
414
+ sx, sy, sz, radii, acol, aflat, order,
415
+ barr, bcol, bflat,
416
+ zoom, ox_s, oy_s,
417
+ l0, l1, l2, f0, f1, f2, hv0, hv1, hv2,
418
+ ambient, fill_w, spec_w, depth_cue, zmin, zspan, n_squares)