원자 오비탈을 셰이더로 볼륨 렌더링하기

원자 구조를 검색하면 매끈한 아령 모양 그림이 나온다. 빨강과 파랑으로 칠해진, 표면이 반들반들한 3D 모델. 그런데 그건 전자의 모양이 아니다. "전자를 발견할 확률이 90%인 지점들을 이어 만든 등확률면"이다. 임의로 정한 90% 선에 껍데기를 씌운 것이라, 실제보다 훨씬 단단해 보인다.

실제 전자는 안개다. 어디에 있다고 말할 수 없고 밀도만 있다. 그러면 안개를 안개로 그려보자는 게 이 글의 내용이다. Figma와 Konva의 캔버스 렌더링 방식 비교에서 WebGL을 2D 에디터에 쓰는 이야기를 했는데, 이번엔 원래 용도에 가까운 쪽이다.

결과물부터 보면 이렇다. 탄소 원자 하나의 전자 여섯 개를 오비탈별로 켜고 끌 수 있다.

탄소 C · Z_eff 5.67 / 3.14

1s² 2s² 2px¹ 2py¹

2px

2py

2pz

드래그 회전 · 휠 확대

안개 감도1.00카메라 거리10.9샘플 수128

전자 6개 표시 / 탄소 C 전체 6개 · 밝기 오비탈별 정규화

라이브러리 없이 WebGL2 프래그먼트 셰이더만으로 렌더링한다. 표면이 없고, 밝기가 곧 확률밀도다.

물리적인 배경이 궁금하면 원자 해부도에 정리해뒀다. 이 글은 구현 쪽이다.

무엇을 그리는 건가

양자역학에서 전자의 상태는 파동함수 ψ로 기술된다. 그리고 ψ를 제곱한 값이 그 지점에서 전자를 발견할 확률밀도가 된다. 우리가 그려야 하는 건 |ψ|²라는 3차원 스칼라 필드다.

여기서 운이 좋은 점이 하나 있다. 수소처럼 전자가 하나인 원자는 파동함수가 닫힌 형태로 딱 떨어진다. 수치해석이 필요 없다는 뜻이다. a₀(보어 반지름)을 1로 두고 정리하면 이렇게 된다.

1s  ψ = 7.616 · e^(−5.67r)
2s  ψ = 0.555 · (2 − 3.14r) · e^(−1.57r)
2px ψ = 1.742 · x · e^(−1.57r)
2py ψ = 1.742 · y · e^(−1.57r)
2pz ψ = 1.742 · z · e^(−1.57r)

지수의 계수는 유효 핵전하다. 탄소를 기준으로 슬레이터 규칙을 적용하면 1s가 5.67, 2s와 2p가 3.14가 나온다. 이 차이 때문에 1s가 2s보다 7배쯤 핵에 가까이 붙는다. 화면에서 보이는 크기 차이가 임의로 정한 값이 아니라 이 숫자에서 나온 결과라는 게 중요하다.

식을 보면 재밌는 게 몇 개 보인다. 2s에는 (2 − 3.14r) 항이 있어서 r = 0.64 지점에서 부호가 뒤집힌다. 그 구면 전체가 확률이 정확히 0인 마디(node)다. 2s는 통짜 공이 아니라 공 속에 공이 든 구조라는 얘기다. 2p는 좌표 성분이 그대로 곱해져 있으니 x = 0 평면 전체가 마디가 된다. 핵 바로 위에서 전자를 발견할 확률이 0인 것이다. 궤도를 도는 공이라고 생각하면 설명할 수 없는 성질이고, 그래서 저 매끈한 아령 그림이 오해를 부른다.

라이브러리를 고르다가 안 쓰기로 했다

처음엔 당연히 three.js 를 쓸 생각이었다. 볼륨 렌더링 예제도 있고 Data3DTexture에 밀도를 채워 레이마칭하는 게 정석이다. React 프로젝트면 React Three Fiber 를 얹으면 되고.

이 분야엔 특화 도구도 있다. 3Dmol.js 는 분자 시각화 전용이라 cube 파일을 읽어 분자 오비탈 등확률면을 그리는 기능이 내장돼 있다. 실제 양자화학 계산 결과를 띄우려면 최단 경로다. Plotly.js 는 isosurface 트레이스에 그리드 데이터만 넘기면 회전까지 공짜로 붙는다. 30분 프로토타입용으로는 이게 제일 빠르다.

그런데 요구사항을 다시 보니 셋 다 과했다. 필요한 건 화면 전체를 덮는 삼각형 하나와 프래그먼트 셰이더 하나뿐이다. 지오메트리도 없고, 조명도 없고, 씬 그래프도 필요 없다. 그리고 파동함수가 닫힌 형태니까 3D 텍스처조차 안 만들어도 된다. 광선 위의 샘플점 좌표를 그 자리에서 식에 넣으면 끝난다.

이렇게 하면 부수 효과로 얻는 게 있다. 격자 해상도라는 개념이 아예 없어서 아무리 확대해도 계단이 생기지 않는다. 텍스처 방식은 64³이든 256³이든 결국 격자에 갇히는데, 해석적으로 계산하면 그 한계가 없다.

구현

전체 파이프라인은 이렇다.

풀스크린 삼각형 → 픽셀별 광선 생성 → 구와 교차 판정
  → 광선 따라 샘플링 → |ψ|² 누적 → 톤매핑

풀스크린 삼각형

사각형 대신 삼각형 하나를 쓴다. 정점 세 개로 화면을 다 덮을 수 있어서 인덱스 버퍼도 필요 없다.

gl.bufferData(
  gl.ARRAY_BUFFER,
  new Float32Array([-1, -1, 3, -1, -1, 3]),
  gl.STATIC_DRAW
);
// ...
gl.drawArrays(gl.TRIANGLES, 0, 3);

좌표가 3까지 나가는 게 핵심이다. NDC 범위 [-1, 1]을 넘어가는 큰 삼각형이라 화면 밖은 알아서 잘린다.

광선 만들기

카메라 행렬을 만들 필요도 없다. yaw, pitch, distance로 카메라 위치를 계산하고 기저 벡터 세 개를 유니폼으로 넘긴 다음, 셰이더에서 방향만 조립해준다.

vec3 rd = normalize(
  uRight * (vNdc.x * uAspect * uTanHalf) +
  uUp    * (vNdc.y * uTanHalf) +
  uFwd
);

여기서 uRight는 cross(forward, worldUp)으로 만든다. 부호를 반대로 잡으면 화면이 좌우로 뒤집힌 채 렌더링되는데, 회전하는 안개라 처음엔 뒤집혔는지도 눈치채기 어렵다. 축을 하나 그려놓고 확인하는 게 낫다.

반지름 7의 구로 자르기

박스 대신 구로 볼륨을 자른다. 코너 처리가 없어서 코드가 짧고, 어차피 원자는 둥글다.

float b = dot(uEye, rd);
float c = dot(uEye, uEye) - R_MAX * R_MAX;
float h = b * b - c;
if (h > 0.0) {
  h = sqrt(h);
  float t0 = max(-b - h, 0.0);
  float t1 = -b + h;
  // t0 ~ t1 구간을 행진한다
}

R_MAX는 7로 뒀다. 지수가 e^(−1.57r)이니 r이 5만 넘어가도 밀도가 무시할 수준이라 더 키울 이유가 없다. 볼륨을 크게 잡으면 같은 샘플 수로 훑는 구간이 길어져서 스텝 간격만 거칠어진다.

누적

표준 emission–absorption 모델이다. 각 스텝에서 밀도만큼 불투명도를 쌓고, 이미 쌓인 만큼 뒤쪽 기여를 깎는다. 앞쪽이 뒤쪽을 가리는 게 자연스럽게 처리된다.

for (int i = 0; i < 512; i++) {
  if (i >= uSteps || alpha > 0.995) break;
  vec3 p = uEye + rd * t;
  float r = length(p);

  float e1 = exp(-5.67 * r);   // 1s
  float e2 = exp(-1.57 * r);   // 2s와 2p가 공유

  float dens = 0.0;
  vec3  emit = vec3(0.0);
  float s, d;

  s = 7.616 * e1;                        d = s * s * uW[0];
  dens += d; emit += d * uCol[0];
  s = 0.5548 * (2.0 - 3.14 * r) * e2;    d = s * s * uW[1];
  dens += d; emit += d * uCol[1];
  s = 1.7423 * p.x * e2;                 d = s * s * uW[2];
  dens += d; emit += d * uCol[2];
  // 2py, 2pz 동일

  if (dens > 1e-9) {
    float a = 1.0 - exp(-dens * uGain * dt * 3.2);
    acc   += (1.0 - alpha) * a * (emit / dens);
    alpha += (1.0 - alpha) * a;
  }
  t += dt;
}

uW는 오비탈별 가중치다. 0을 넣으면 그 오비탈이 꺼지고, 값을 바꾸면 밝기가 조절된다. 조건 분기 대신 가중치를 곱하는 쪽이 셰이더에서 다루기 편하다.

emit / dens로 나누는 부분이 색 결정이다. 여러 오비탈이 겹친 지점에서는 밀도로 가중평균한 색이 나온다. 겹치는 영역이 회색으로 수렴하는데, 이건 버그가 아니라 정직한 결과다. 실제로 그 지점에는 여러 오비탈이 다 있다.

성능

프래그먼트 셰이더가 픽셀마다 100번 넘게 도는 구조라 스텝당 비용이 그대로 곱해진다. 세 가지로 잡았다.

지수 계산을 공유한 게 제일 컸다. 2s와 2p가 모두 e^(−1.57r)을 쓰니까 한 번만 계산해서 넷이 나눠 쓴다. 오비탈마다 따로 exp()를 부르면 스텝당 5번인데 2번으로 줄어든다.

그리고 dirty 플래그를 뒀다. 자동 회전이 꺼져 있고 사용자가 아무것도 안 건드리면 그리지 않는다. requestAnimationFrame은 계속 돌지만 drawArrays는 호출되지 않는다.

마지막으로 devicePixelRatio를 1.25로 캡했다. 레티나에서 2배로 그리면 픽셀 수가 4배가 되는데, 안개는 원래 흐릿해서 해상도를 다 쓸 이유가 없다.

샘플 수는 슬라이더로 빼놨다. 기기 성능을 미리 알 수 없으니 사용자가 내리게 하는 편이 안전하다.

마크다운 글 안에 넣기

이 블로그는 마크다운을 react-markdown 으로 렌더링한다. rehype-raw가 붙어 있어서 raw HTML은 통과하지만, 그 안의 <script>는 실행되지 않는다. React가 만든 엘리먼트에 인라인 스크립트를 넣어도 브라우저가 실행하지 않기 때문이다. 그래서 HTML을 붙여넣는 방식으로는 인터랙티브 뷰어를 넣을 수 없다.

대신 이미 있던 패턴을 재사용했다. 이 블로그는 ```mermaid 펜스를 컴포넌트로 바꿔치기하고 있었으니, 같은 자리에 언어 하나를 추가하면 된다.

code: ({ className, children, ...props }) => {
  const language = /language-(\w+)/.exec(className || "")?.[1] ?? "";

  if (language === "orbital") {
    const opts = parseOrbitalOptions(String(children));
    return <OrbitalViewer {...opts} />;
  }
  // ...
}

펜스 내용은 JSON으로 파싱해서 props로 넘긴다. 어떤 원자를 그릴지, 어떤 오비탈을 켠 상태로 시작할지, 높이와 캡션을 글마다 다르게 줄 수 있다.

한 가지 더 손봐야 했다. pre 핸들러가 코드 블록마다 복사 버튼과 어두운 배경을 씌우고 있어서, 뷰어에도 그게 붙는다. 자식 code 엘리먼트의 className을 보고 우회하게 만들었다.

if (codeElement?.props?.className?.includes("language-orbital")) {
  return <>{children}</>;
}

삽질 세 가지

readPixels가 전부 0으로 나온다

렌더링이 되는지 확인하려고 픽셀을 읽었는데 전부 0이었다. 셰이더는 컴파일도 링크도 통과했고 콘솔에 에러가 없는데 검은 화면이 나오는 것처럼 보였다.

원인은 WebGL의 drawing buffer가 컴포지팅 후에 비워지기 때문이다. preserveDrawingBuffer: true를 주지 않으면 프레임이 화면에 합성된 다음 버퍼가 클리어된다. 그러니 requestAnimationFrame 콜백 밖에서 읽으면 항상 빈 버퍼를 받는다.

읽는 시점을 프레임 안으로 옮기면 해결된다.

requestAnimationFrame(() => {
  gl.readPixels(0, 0, w, h, gl.RGBA, gl.UNSIGNED_BYTE, buf);
});

한 가지 더 걸렸다. dirty 플래그를 넣어놨으니 자동 회전을 끈 상태에서는 그 프레임에 drawArrays가 아예 안 불린다. 그리지 않은 프레임에서 읽으면 당연히 0이다. 최적화와 검증이 부딪히는 지점인데, 읽기 직전에 상태를 한 번 건드려서 다시 그리게 하면 된다.

톤매핑이 색을 흰색으로 날려버린다

HDR로 누적한 값을 채널별로 압축하고 감마를 씌우면 밝은 부분의 색이 죄다 흰색으로 수렴한다. 감마 곡선이 큰 값들을 서로 가깝게 끌어당기니까 채널 비율이 무너지는 것이다. 코어가 하얗게 타는 건 그럴싸해 보이지만, 위상 부호를 색으로 구분하는 기능을 넣어놨는데 그게 안 보이면 의미가 없다.

두 가지를 바꿨다. 압축을 채널별로 하지 않고 최대 성분 기준으로 해서 채널 비율을 유지하고, 감마 뒤에 채도를 다시 끌어냈다.

float m = max(max(col.r, col.g), col.b);
if (m > 1e-6) col *= (m / (m + 0.85)) / m;
col = pow(col, vec3(0.4545));

float lum = dot(col, vec3(0.299, 0.587, 0.114));
col = clamp(mix(vec3(lum), col, 1.45), 0.0, 1.0);

렌더 결과를 직접 측정해서 확인했다. 단일 오비탈에 위상 색을 켠 상태에서 채도 0.35를 넘는 픽셀 비율이 34%에서 84%로 올라갔다.

이런 건 눈으로 보면서 감으로 맞추기 어렵다. 픽셀을 읽어 통계를 내면 숫자로 비교할 수 있어서 훨씬 빠르다.

물리 설명이 투영에서 틀렸다

2px의 마디는 x = 0 평면 전체다. 그래서 "핵 위치는 완전히 검게 보인다"고 설명을 적어놨는데, 실제로 화면 중앙 픽셀을 읽어보니 밝았다.

당연한 결과였다. 볼륨 렌더링은 광선이 지나온 경로를 다 누적한다. 카메라가 x축과 어긋난 방향에 있으면 화면 중앙을 지나는 광선도 x가 0이 아닌 지점들을 통과하니까 밀도가 쌓인다. 마디가 검게 보이는 건 광선 전체가 그 평면 위에 놓일 때, 즉 아령 축을 화면과 나란히 뒀을 때뿐이다.

3D 필드의 성질과 그 필드를 2D로 투영한 결과는 다르다. 필드에서 참인 문장을 그림 설명에 그대로 옮기면 틀릴 수 있다. 설명 문구를 "축을 정면으로 보면 두 로브 사이에 검은 띠가 드러난다"로 고쳤다.

결과

정지 그림으로는 안 보이던 것들이 몇 개 드러난다. 위 뷰어에서 직접 확인할 수 있다.

2s만 켜고 감도를 올리면 밝은 안쪽 덩어리와 흐린 바깥 껍질 사이에 어두운 고리가 뜬다. 마디 구면을 옆에서 본 단면이다. 2p 세 방향을 다 켜면 방향성이 사라지고 전체가 다시 둥글어진다. 네온이나 아르곤처럼 p가 꽉 찬 원소가 반응하지 않는 이유가 그림으로 보인다.

제일 인상적인 건 밝기 정규화를 끄는 쪽이다. 기본값은 오비탈마다 밝기를 맞춰놨는데, 실제 밀도 비율로 바꾸면 1s만 남고 나머지가 거의 사라진다. 전자 6개 중 2개가 핵에 딱 붙어 있고 화학 반응에 참여하는 건 바깥의 흐릿한 쪽이라는 게 눈에 들어온다. 교과서 그림에서는 이 비율이 절대 안 보인다.

정리하면 볼륨 렌더링은 등확률면보다 코드가 짧고 정직하다. 90%라는 임의의 선을 고르지 않아도 되고, 마칭 큐브도 필요 없다. 데이터를 격자에 굽지 않고 식을 그대로 쓸 수 있는 상황이면 더 그렇다.

참고

Hydrogen-like atom - Wikipedia — 수소형 파동함수의 해석적 형태
Slater's rules - Wikipedia — 유효 핵전하 계산
Volume ray casting - Wikipedia — emission–absorption 누적 모델
WebGL2RenderingContext - MDN — WebGL2 API 레퍼런스
three.js examples — 3D 텍스처 기반 볼륨 렌더링 예제
3Dmol.js — 분자 오비탈 시각화 전용 라이브러리
Plotly.js 3D Isosurface — 등확률면을 빠르게 띄우는 방법