File: water.py

package info (click to toggle)
python-moderngl 5.12.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 4,700 kB
  • sloc: python: 15,758; cpp: 14,665; makefile: 14
file content (182 lines) | stat: -rw-r--r-- 5,634 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
# WebGL Water
# https://madebyevan.com/webgl-water/
# Copyright 2011 Evan Wallace
# Released under the MIT license

import moderngl
from light_gl import *


class Water:
    _vert_shader = """
    varying vec2 coord;
    void main() {
        coord = gl_Vertex.xy * 0.5 + 0.5;
        gl_Position = vec4(gl_Vertex.xyz, 1.0);
    }
    """

    _drop_frag_shader = """
    const float PI = 3.141592653589793;
    uniform sampler2D texture;
    uniform vec2 center;
    uniform float radius;
    uniform float strength;
    varying vec2 coord;
    void main() {
        /* get vertex info */
        vec4 info = texture2D(texture, coord);
        /* add the drop to the height */
        float drop = max(0.0, 1.0 - length(center * 0.5 + 0.5 - coord) / radius);
        drop = 0.5 - cos(drop * PI) * 0.5;
        info.r += drop * strength;
        gl_FragColor = info;
    }
    """

    _update_frag_shader = """
    uniform sampler2D texture;
    uniform vec2 delta;
    varying vec2 coord;

    vec4 water_sample(vec2 coord, vec4 info) {
        // vec2 center = vec2(0.5, 0.4);
        // float radius = 0.10;  
        // if (length(coord.xy - center) < radius) {
        //     return info;
        // }
        return texture2D(texture, coord);
    }

    void main() {
        /* get vertex info */
        vec4 info = water_sample(coord, vec4(0.5));
        
        /* calculate average neighbor height */
        vec2 dx = vec2(delta.x, 0.0);
        vec2 dy = vec2(0.0, delta.y);
        float average = (
            water_sample(coord - dx, info).r +
            water_sample(coord - dy, info).r +
            water_sample(coord + dx, info).r +
            water_sample(coord + dy, info).r
        ) * 0.25;
        
        /* change the velocity to move toward the average */
        info.g += (average - info.r) * 2.0;
        
        /* attenuate the velocity a little so waves do not last forever */
        info.g *= 0.995;
        
        /* move the vertex along the velocity */
        info.r += info.g;

        gl_FragColor = vec4(info.rg, 0.0, 0.0);
    }
    """

    _normal_frag_shader = """
    uniform sampler2D texture;
    uniform vec2 delta;
    varying vec2 coord;
    void main() {
        /* get vertex info */
        vec4 info = texture2D(texture, coord);
        
        /* update the normal */
        vec3 dx = vec3(delta.x, texture2D(texture, vec2(coord.x + delta.x, coord.y)).r - info.r, 0.0);
        vec3 dy = vec3(0.0, texture2D(texture, vec2(coord.x, coord.y + delta.y)).r - info.r, delta.y);
        info.ba = normalize(cross(dy, dx)).xz;
        
        gl_FragColor = info;
    }
    """

    _sphere_frag_shader = """
    uniform sampler2D texture;
    uniform vec3 oldCenter;
    uniform vec3 newCenter;
    uniform float radius;
    varying vec2 coord;
    
    float volumeInSphere(vec3 center) {
      vec3 toCenter = vec3(coord.x * 2.0 - 1.0, 0.0, coord.y * 2.0 - 1.0) - center;
      float t = length(toCenter) / radius;
      float dy = exp(-pow(t * 1.5, 6.0));
      float ymin = min(0.0, center.y - dy);
      float ymax = min(max(0.0, center.y + dy), ymin + 2.0 * dy);
      return (ymax - ymin) * 0.1;
    }
    
    void main() {
      /* get vertex info */
      vec4 info = texture2D(texture, coord);
      
      /* add the old volume */
      info.r += volumeInSphere(oldCenter);
      
      /* subtract the new volume */
      info.r -= volumeInSphere(newCenter);
      
      gl_FragColor = info;
    }
    """

    def __init__(self, texture_size=(256, 256), ctx: moderngl.Context = None) -> None:
        self.ctx = ctx or None
        self.panel = MeshBuilder.panel().build(self.ctx)
        self.texture_a = RawTexture(
            texture_size, pixel_fmt="RGBA", ctx=self.ctx, dtype="f4")
        self.texture_b = RawTexture(
            texture_size, pixel_fmt="RGBA", ctx=self.ctx, dtype="f4")
        self.delta = glm.vec2(
            1.0 / texture_size[0], 1.0 / texture_size[1])
        self.drop_shader = Shader(
            Water._vert_shader, Water._drop_frag_shader, self.ctx)
        self.update_shader = Shader(
            Water._vert_shader, Water._update_frag_shader, self.ctx)
        self.normal_shader = Shader(
            Water._vert_shader, Water._normal_frag_shader, self.ctx)
        self.sphere_shader = Shader(
            Water._vert_shader, Water._sphere_frag_shader, self.ctx)

    def swap(self):
        temp_texture = self.texture_b
        self.texture_b = self.texture_a
        self.texture_a = temp_texture

    def add_drop(self, x: float, y: float, radius: float, strength: float):
        self.texture_b.draw_to(self.ctx)
        self.texture_a.use()
        self.drop_shader.draw_mesh(self.panel, unifroms={
            "center": glm.vec2(x, y),
            "radius": radius,
            "strength": strength
        })
        self.swap()

    def move_sphere(self, old_center, new_center, radius):
        self.texture_b.draw_to(self.ctx)
        self.texture_a.use()
        self.sphere_shader.draw_mesh(self.panel, unifroms={
            "oldCenter": old_center,
            "newCenter": new_center,
            "radius": radius
        })
        self.swap()

    def step_simulation(self):
        self.texture_b.draw_to(self.ctx)
        self.texture_a.use()
        self.update_shader.draw_mesh(self.panel, unifroms={
            "delta": self.delta
        })
        self.swap()
    
    def update_normals(self):
        self.texture_b.draw_to(self.ctx)
        self.texture_a.use()
        self.normal_shader.draw_mesh(self.panel, unifroms={
            "delta": self.delta
        })
        self.swap()