sun_position/draw.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import bpy
   4 import math
   5 import gpu
   6 from gpu_extras.batch import batch_for_shader
   7 from mathutils import Vector
   8
   9 from .sun_calc import calc_surface, calc_analemma
  10
  11
  12 if bpy.app.background:  # ignore north line in background mode
  13     def north_update(self, context):
  14         pass
  15     def surface_update(self, context):
  16         pass
  17     def analemmas_update(self, context):
  18         pass
  19 else:
  20     # North line
  21
  22     shader_interface = gpu.types.GPUStageInterfaceInfo("my_interface")
  23     shader_interface.flat('VEC2', "v_StartPos")
  24     shader_interface.smooth('VEC4', "v_VertPos")
  25
  26     shader_info = gpu.types.GPUShaderCreateInfo()
  27     shader_info.push_constant('MAT4', "u_ViewProjectionMatrix")
  28     shader_info.push_constant('VEC4', "u_Color")
  29     shader_info.push_constant('VEC2', "u_Resolution")
  30     shader_info.vertex_in(0, 'VEC3', "position")
  31     shader_info.vertex_out(shader_interface)
  32
  33     shader_info.vertex_source(
  34         '''
  35         void main()
  36         {
  37             vec4 pos    = u_ViewProjectionMatrix * vec4(position, 1.0f);
  38             gl_Position = pos;
  39             v_StartPos    = (pos / pos.w).xy;
  40             v_VertPos     = pos;
  41         }
  42         '''
  43     )
  44
  45     shader_info.fragment_out(0, 'VEC4', "FragColor")
  46     shader_info.fragment_source(
  47         '''
  48         void main()
  49         {
  50             vec4 vertPos_2d = v_VertPos / v_VertPos.w;
  51             vec2 dir  = (vertPos_2d.xy - v_StartPos.xy) * u_Resolution;
  52             float dist = length(dir);
  53
  54             if (step(sin(dist / 5.0f), 0.0) == 1) discard;
  55
  56             FragColor = u_Color;
  57         }
  58         '''
  59     )
  60
  61     shader = gpu.shader.create_from_info(shader_info)
  62     del shader_info
  63     del shader_interface
  64
  65     def north_draw():
  66         """
  67         Set up the compass needle using the current north offset angle
  68         less 90 degrees.  This forces the unit circle to begin at the
  69         12 O'clock instead of 3 O'clock position.
  70         """
  71         sun_props = bpy.context.scene.sun_pos_properties
  72
  73         color = (0.2, 0.6, 1.0, 0.7)
  74         radius = 100
  75         angle = -(sun_props.north_offset - math.pi / 2)
  76         x = math.cos(angle) * radius
  77         y = math.sin(angle) * radius
  78         coords = Vector((x, y, 0)), Vector((0, 0, 0))
  79         batch = batch_for_shader(shader, 'LINE_STRIP', {"position": coords})
  80
  81         matrix = bpy.context.region_data.perspective_matrix
  82         shader.uniform_float("u_ViewProjectionMatrix", matrix)
  83         shader.uniform_float("u_Resolution", (bpy.context.region.width,
  84                                               bpy.context.region.height))
  85         shader.uniform_float("u_Color", color)
  86         width = gpu.state.line_width_get()
  87         gpu.state.line_width_set(2.0)
  88         batch.draw(shader)
  89         gpu.state.line_width_set(width)
  90
  91     _north_handle = None
  92
  93     def north_update(self, context):
  94         global _north_handle
  95         sun_props = context.scene.sun_pos_properties
  96         addon_prefs = context.preferences.addons[__package__].preferences
  97
  98         if addon_prefs.show_overlays and sun_props.show_north:
  99             _north_handle = bpy.types.SpaceView3D.draw_handler_add(north_draw, (), 'WINDOW', 'POST_VIEW')
 100         elif _north_handle is not None:
 101             bpy.types.SpaceView3D.draw_handler_remove(_north_handle, 'WINDOW')
 102             _north_handle = None
 103
 104     # Analemmas
 105
 106     def analemmas_draw(batch, shader):
 107         shader.uniform_float("color", (1, 0, 0, 1))
 108         batch.draw(shader)
 109
 110     _analemmas_handle = None
 111
 112     def analemmas_update(self, context):
 113         global _analemmas_handle
 114         sun_props = context.scene.sun_pos_properties
 115         addon_prefs = context.preferences.addons[__package__].preferences
 116
 117         if addon_prefs.show_overlays and sun_props.show_analemmas:
 118             coords = []
 119             indices = []
 120             coord_offset = 0
 121             for h in range(24):
 122                 analemma_verts = calc_analemma(context, h)
 123                 coords.extend(analemma_verts)
 124                 for i in range(len(analemma_verts) - 1):
 125                     indices.append((coord_offset + i,
 126                                     coord_offset + i+1))
 127                 coord_offset += len(analemma_verts)
 128
 129             shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR')
 130             batch = batch_for_shader(shader, 'LINES',
 131                                     {"pos": coords}, indices=indices)
 132
 133             if _analemmas_handle is not None:
 134                 bpy.types.SpaceView3D.draw_handler_remove(_analemmas_handle, 'WINDOW')
 135             _analemmas_handle = bpy.types.SpaceView3D.draw_handler_add(
 136                 analemmas_draw, (batch, shader), 'WINDOW', 'POST_VIEW')
 137         elif _analemmas_handle is not None:
 138             bpy.types.SpaceView3D.draw_handler_remove(_analemmas_handle, 'WINDOW')
 139             _analemmas_handle = None
 140
 141     # Surface
 142
 143     def surface_draw(batch, shader):
 144         blend = gpu.state.blend_get()
 145         gpu.state.blend_set("ALPHA")
 146         shader.uniform_float("color", (.8, .6, 0, 0.2))
 147         batch.draw(shader)
 148         gpu.state.blend_set(blend)
 149
 150     _surface_handle = None
 151
 152     def surface_update(self, context):
 153         global _surface_handle
 154         sun_props = context.scene.sun_pos_properties
 155         addon_prefs = context.preferences.addons[__package__].preferences
 156
 157         if addon_prefs.show_overlays and sun_props.show_surface:
 158             coords = calc_surface(context)
 159             shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR')
 160             batch = batch_for_shader(shader, 'TRIS', {"pos": coords})
 161
 162             if _surface_handle is not None:
 163                 bpy.types.SpaceView3D.draw_handler_remove(_surface_handle, 'WINDOW')
 164             _surface_handle = bpy.types.SpaceView3D.draw_handler_add(
 165                 surface_draw, (batch, shader), 'WINDOW', 'POST_VIEW')
 166         elif _surface_handle is not None:
 167             bpy.types.SpaceView3D.draw_handler_remove(_surface_handle, 'WINDOW')
 168             _surface_handle = None